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, PIPE 

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 # For Python 3.6 and possibly older 

44 stdout=PIPE, stderr=PIPE, universal_newlines=True, 

45 # For Python 3.7+ the next line is equivalent 

46 # capture_output=True, text=True, 

47 check=True) 

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

49 

50 

51def check_imports(expression: str, *, 

52 forbidden_modules: List[str] = [], 

53 max_module_count: Optional[int] = None, 

54 max_nonstdlib_module_count: Optional[int] = None, 

55 do_print: bool = False) -> None: 

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

57 

58 Parameters 

59 ---------- 

60 expression 

61 Python expression 

62 forbidden_modules 

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

64 max_module_count 

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

66 max_nonstdlib_module_count 

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

68 do_print: 

69 Print loaded modules if set. 

70 """ 

71 modules = exec_and_check_modules(expression) 

72 

73 if do_print: 

74 print('all modules:') 

75 pprint(sorted(modules)) 

76 

77 for module_pattern in forbidden_modules: 

78 r = re.compile(module_pattern) 

79 for module in modules: 

80 assert not r.fullmatch(module), \ 

81 f'{module} was imported' 

82 

83 if max_nonstdlib_module_count is not None: 

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

85 

86 nonstdlib_modules = [] 

87 for module in modules: 

88 if module.split('.')[0] in sys.stdlib_module_names: # type: ignore 

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))