Coverage for ase / utils / checkimports.py: 82.86%

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 

19 

20 

21def exec_and_check_modules(expression: str) -> set[str]: 

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

23 

24 Parameters 

25 ---------- 

26 expression 

27 Python expression 

28 

29 Returns 

30 ------- 

31 Set of module names. 

32 """ 

33 # Take null outside command to avoid 

34 # `import os` before expression 

35 null = os.devnull 

36 command = ( 

37 'import sys;' 

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

39 f' {expression};' 

40 ' sys.stdout = stdout;' 

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

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

43 ) 

44 proc = run( 

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

46 capture_output=True, 

47 universal_newlines=True, 

48 check=True, 

49 ) 

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

51 

52 

53def check_imports( 

54 expression: str, 

55 *, 

56 forbidden_modules: list[str] = [], 

57 max_module_count: int | None = None, 

58 max_nonstdlib_module_count: int | None = None, 

59 do_print: bool = False, 

60) -> None: 

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

62 

63 Parameters 

64 ---------- 

65 expression 

66 Python expression 

67 forbidden_modules 

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

69 max_module_count 

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

71 max_nonstdlib_module_count 

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

73 do_print: 

74 Print loaded modules if set. 

75 """ 

76 modules = exec_and_check_modules(expression) 

77 

78 if do_print: 

79 print('all modules:') 

80 pprint(sorted(modules)) 

81 

82 for module_pattern in forbidden_modules: 

83 r = re.compile(module_pattern) 

84 for module in modules: 

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

86 

87 if max_nonstdlib_module_count is not None: 

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

89 

90 nonstdlib_modules = [] 

91 for module in modules: 

92 if ( 

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

94 ): 

95 continue 

96 nonstdlib_modules.append(module) 

97 

98 if do_print: 

99 print('nonstdlib modules:') 

100 pprint(sorted(nonstdlib_modules)) 

101 

102 module_count = len(nonstdlib_modules) 

103 assert module_count <= max_nonstdlib_module_count, ( 

104 'too many nonstdlib modules loaded:' 

105 f' {module_count}/{max_nonstdlib_module_count}' 

106 ) 

107 

108 if max_module_count is not None: 

109 module_count = len(modules) 

110 assert module_count <= max_module_count, ( 

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

112 ) 

113 

114 

115if __name__ == '__main__': 

116 import argparse 

117 

118 parser = argparse.ArgumentParser() 

119 parser.add_argument('expression') 

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

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

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

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

124 args = parser.parse_args() 

125 

126 check_imports(**vars(args))