Coverage for /builds/ase/ase/ase/utils/checkimports.py: 67.39%

1"""Utility for checking Python module imports triggered by any code snippet. 

2 

3This module was developed to monitor the import footprint of the ase CLI 

4command: The CLI command can become unnecessarily slow and unresponsive 

5if too many modules are imported even before the CLI is launched or 

6it is known what modules will be actually needed. 

7See https://gitlab.com/ase/ase/-/issues/1124 for more discussion. 

8 

9The utility here is general, so it can be used for checking and 

10monitoring other code snippets too. 

11""" 

12 

13import json 

14import os 

15import re 

16import sys 

17from pprint import pprint 

18from subprocess import run 

19from typing import List, Optional, Set 

20 

21 

22def exec_and_check_modules(expression: str) -> Set[str]: 

23 """Return modules loaded by the execution of a Python expression. 

24 

25 Parameters 

26 ---------- 

27 expression 

28 Python expression 

29 

30 Returns 

31 ------- 

32 Set of module names. 

33 """ 

34 # Take null outside command to avoid 

35 # `import os` before expression 

36 null = os.devnull 

37 command = ( 

38 'import sys;' 

39 f" stdout = sys.stdout; sys.stdout = open({null!r}, 'w');" 

40 f' {expression};' 

41 ' sys.stdout = stdout;' 

42 ' modules = list(sys.modules);' 

43 ' import json; print(json.dumps(modules))' 

44 ) 

45 proc = run( 

46 [sys.executable, '-c', command], 

47 capture_output=True, 

48 universal_newlines=True, 

49 check=True, 

50 ) 

51 return set(json.loads(proc.stdout)) 

52 

53 

54def check_imports( 

55 expression: str, 

56 *, 

57 forbidden_modules: List[str] = [], 

58 max_module_count: Optional[int] = None, 

59 max_nonstdlib_module_count: Optional[int] = None, 

60 do_print: bool = False, 

61) -> None: 

62 """Check modules imported by the execution of a Python expression. 

63 

64 Parameters 

65 ---------- 

66 expression 

67 Python expression 

68 forbidden_modules 

69 Throws an error if any module in this list was loaded. 

70 max_module_count 

71 Throws an error if the number of modules exceeds this value. 

72 max_nonstdlib_module_count 

73 Throws an error if the number of non-stdlib modules exceeds this value. 

74 do_print: 

75 Print loaded modules if set. 

76 """ 

77 modules = exec_and_check_modules(expression) 

78 

79 if do_print: 

80 print('all modules:') 

81 pprint(sorted(modules)) 

82 

83 for module_pattern in forbidden_modules: 

84 r = re.compile(module_pattern) 

85 for module in modules: 

86 assert not r.fullmatch(module), f'{module} was imported' 

87 

88 if max_nonstdlib_module_count is not None: 

89 assert sys.version_info >= (3, 10), 'Python 3.10+ required' 

90 

91 nonstdlib_modules = [] 

92 for module in modules: 

93 if ( 

94 module.split('.')[0] in sys.stdlib_module_names # type: ignore[attr-defined] 

95 ): 

96 continue 

97 nonstdlib_modules.append(module) 

98 

99 if do_print: 

100 print('nonstdlib modules:') 

101 pprint(sorted(nonstdlib_modules)) 

102 

103 module_count = len(nonstdlib_modules) 

104 assert module_count <= max_nonstdlib_module_count, ( 

105 'too many nonstdlib modules loaded:' 

106 f' {module_count}/{max_nonstdlib_module_count}' 

107 ) 

108 

109 if max_module_count is not None: 

110 module_count = len(modules) 

111 assert module_count <= max_module_count, ( 

112 f'too many modules loaded: {module_count}/{max_module_count}' 

113 ) 

114 

115 

116if __name__ == '__main__': 

117 import argparse 

118 

119 parser = argparse.ArgumentParser() 

120 parser.add_argument('expression') 

121 parser.add_argument('--forbidden_modules', nargs='+', default=[]) 

122 parser.add_argument('--max_module_count', type=int, default=None) 

123 parser.add_argument('--max_nonstdlib_module_count', type=int, default=None) 

124 parser.add_argument('--do_print', action='store_true') 

125 args = parser.parse_args() 

126 

127 check_imports(**vars(args))