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

1# fmt: off 

2 

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

4 

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

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

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

8it is known what modules will be actually needed. 

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

10 

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

12monitoring other code snippets too. 

13""" 

14import json 

15import os 

16import re 

17import sys 

18from pprint import pprint 

19from subprocess import run 

20from typing import List, Optional, Set 

21 

22 

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

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

25 

26 Parameters 

27 ---------- 

28 expression 

29 Python expression 

30 

31 Returns 

32 ------- 

33 Set of module names. 

34 """ 

35 # Take null outside command to avoid 

36 # `import os` before expression 

37 null = os.devnull 

38 command = ("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 proc = run([sys.executable, '-c', command], 

45 capture_output=True, universal_newlines=True, 

46 check=True) 

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

48 

49 

50def check_imports(expression: str, *, 

51 forbidden_modules: List[str] = [], 

52 max_module_count: Optional[int] = None, 

53 max_nonstdlib_module_count: Optional[int] = None, 

54 do_print: bool = False) -> None: 

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

56 

57 Parameters 

58 ---------- 

59 expression 

60 Python expression 

61 forbidden_modules 

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

63 max_module_count 

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

65 max_nonstdlib_module_count 

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

67 do_print: 

68 Print loaded modules if set. 

69 """ 

70 modules = exec_and_check_modules(expression) 

71 

72 if do_print: 

73 print('all modules:') 

74 pprint(sorted(modules)) 

75 

76 for module_pattern in forbidden_modules: 

77 r = re.compile(module_pattern) 

78 for module in modules: 

79 assert not r.fullmatch(module), \ 

80 f'{module} was imported' 

81 

82 if max_nonstdlib_module_count is not None: 

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

84 

85 nonstdlib_modules = [] 

86 for module in modules: 

87 if ( 

88 module.split('.')[0] 

89 in sys.stdlib_module_names # type: ignore[attr-defined] 

90 ): 

91 continue 

92 nonstdlib_modules.append(module) 

93 

94 if do_print: 

95 print('nonstdlib modules:') 

96 pprint(sorted(nonstdlib_modules)) 

97 

98 module_count = len(nonstdlib_modules) 

99 assert module_count <= max_nonstdlib_module_count, ( 

100 'too many nonstdlib modules loaded:' 

101 f' {module_count}/{max_nonstdlib_module_count}' 

102 ) 

103 

104 if max_module_count is not None: 

105 module_count = len(modules) 

106 assert module_count <= max_module_count, \ 

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

108 

109 

110if __name__ == '__main__': 

111 import argparse 

112 

113 parser = argparse.ArgumentParser() 

114 parser.add_argument('expression') 

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

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

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

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

119 args = parser.parse_args() 

120 

121 check_imports(**vars(args))