Coverage for /builds/kinetik161/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""" 

12import json 

13import os 

14import re 

15import sys 

16from pprint import pprint 

17from subprocess import run 

18from typing import List, Optional, Set 

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 = ("import sys;" 

37 f" stdout = sys.stdout; sys.stdout = open({repr(null)}, 'w');" 

38 f" {expression};" 

39 " sys.stdout = stdout;" 

40 " modules = list(sys.modules);" 

41 " import json; print(json.dumps(modules))") 

42 proc = run([sys.executable, '-c', command], 

43 capture_output=True, universal_newlines=True, 

44 check=True) 

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

46 

47 

48def check_imports(expression: str, *, 

49 forbidden_modules: List[str] = [], 

50 max_module_count: Optional[int] = None, 

51 max_nonstdlib_module_count: Optional[int] = None, 

52 do_print: bool = False) -> None: 

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

54 

55 Parameters 

56 ---------- 

57 expression 

58 Python expression 

59 forbidden_modules 

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

61 max_module_count 

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

63 max_nonstdlib_module_count 

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

65 do_print: 

66 Print loaded modules if set. 

67 """ 

68 modules = exec_and_check_modules(expression) 

69 

70 if do_print: 

71 print('all modules:') 

72 pprint(sorted(modules)) 

73 

74 for module_pattern in forbidden_modules: 

75 r = re.compile(module_pattern) 

76 for module in modules: 

77 assert not r.fullmatch(module), \ 

78 f'{module} was imported' 

79 

80 if max_nonstdlib_module_count is not None: 

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

82 

83 nonstdlib_modules = [] 

84 for module in modules: 

85 if ( 

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

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

88 ): 

89 continue 

90 nonstdlib_modules.append(module) 

91 

92 if do_print: 

93 print('nonstdlib modules:') 

94 pprint(sorted(nonstdlib_modules)) 

95 

96 module_count = len(nonstdlib_modules) 

97 assert module_count <= max_nonstdlib_module_count, ( 

98 'too many nonstdlib modules loaded:' 

99 f' {module_count}/{max_nonstdlib_module_count}' 

100 ) 

101 

102 if max_module_count is not None: 

103 module_count = len(modules) 

104 assert module_count <= max_module_count, \ 

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

106 

107 

108if __name__ == '__main__': 

109 import argparse 

110 

111 parser = argparse.ArgumentParser() 

112 parser.add_argument('expression') 

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

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

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

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

117 args = parser.parse_args() 

118 

119 check_imports(**vars(args))