Coverage for src/lazy_imports_lite/_loader.py: 100%

85 statements  

« prev     ^ index     » next       coverage.py v7.4.1, created at 2024-02-12 09:19 +0100

1import ast 

2import importlib.abc 

3import importlib.machinery 

4import importlib.metadata 

5import os 

6import sys 

7import types 

8 

9from ._hooks import LazyObject 

10from ._transformer import TransformModuleImports 

11 

12 

13class LazyModule(types.ModuleType): 

14 def __getattribute__(self, name): 

15 value = super().__getattribute__(name) 

16 if isinstance(value, LazyObject): 

17 return value._lazy_value 

18 return value 

19 

20 def __setattr__(self, name, value): 

21 try: 

22 current_value = super().__getattribute__(name) 

23 except: 

24 super().__setattr__(name, value) 

25 else: 

26 if isinstance(current_value, LazyObject): 

27 current_value._lazy_value = value 

28 else: 

29 super().__setattr__(name, value) 

30 

31 

32enabled_packages = set() 

33 

34 

35def scan_distributions(): 

36 global enabled_packages 

37 for dist in importlib.metadata.distributions(): 

38 metadata = dist.metadata 

39 

40 if metadata is None: 

41 continue # pragma: no cover 

42 

43 if metadata["Keywords"] is None: 

44 continue 

45 

46 keywords = metadata["Keywords"].split(",") 

47 if "lazy-imports-lite-enabled" in keywords: 

48 for pkg in _top_level_declared(dist) or _top_level_inferred(dist): 

49 enabled_packages.add(pkg) 

50 

51 

52def _top_level_declared(dist): 

53 return (dist.read_text("top_level.txt") or "").split() 

54 

55 

56def _top_level_inferred(dist): 

57 files = dist.files 

58 if files is None: 

59 return {} # pragma: no cover 

60 

61 parts = { 

62 f.parts[:-1] if len(f.parts) > 1 else f.with_suffix("").name 

63 for f in files 

64 if f.suffix == ".py" 

65 } 

66 

67 is_namespace = min(len(p) for p in parts) == 2 

68 

69 if is_namespace: 

70 return {".".join(p) for p in parts if len(p) == 2} 

71 else: 

72 return {".".join(p) for p in parts if len(p) == 1} 

73 

74 

75class LazyLoader(importlib.abc.Loader, importlib.machinery.PathFinder): 

76 def find_spec(self, fullname, path=None, target=None): 

77 if fullname.startswith("encodings."): 

78 # fix wired windows bug 

79 return None 

80 

81 if "LAZY_IMPORTS_LITE_DISABLE" in os.environ: 

82 return None 

83 

84 spec = super().find_spec(fullname, path, target) 

85 

86 if spec is None: 

87 return None 

88 

89 if spec.origin is None: 

90 return None # pragma: no cover 

91 

92 name = spec.name.split(".")[0] 

93 namespace_name = ".".join(spec.name.split(".")[:2]) 

94 

95 if ( 

96 name in enabled_packages or namespace_name in enabled_packages 

97 ) and spec.origin.endswith(".py"): 

98 

99 origin: str = spec.origin 

100 with open(origin) as f: 

101 mod_raw = f.read() 

102 mod_ast = ast.parse(mod_raw, origin, "exec") 

103 for node in ast.walk(mod_ast): 

104 if isinstance(node,ast.Call) and isinstance(node.func,ast.Name) and node.func.id in ("eval","exec"): 

105 return None 

106 spec.mod_ast=mod_ast 

107 spec.loader = self 

108 return spec 

109 

110 return None 

111 

112 def create_module(self, spec): 

113 return LazyModule(spec.name) 

114 

115 def exec_module(self, module): 

116 origin: str = module.__spec__.origin 

117 

118 mod_ast = module.__spec__.mod_ast 

119 del module.__spec__.mod_ast 

120 

121 transformer = TransformModuleImports() 

122 new_ast = transformer.visit(mod_ast) 

123 

124 ast.fix_missing_locations(new_ast) 

125 mod_code = compile(new_ast, origin, "exec") 

126 exec(mod_code, module.__dict__) 

127 del module.__dict__["__lazy_imports_lite__"] 

128 del module.__dict__["globals"] 

129 

130 

131def setup(): 

132 scan_distributions() 

133 

134 if not any(isinstance(m, LazyLoader) for m in sys.meta_path): 

135 sys.meta_path.insert(0, LazyLoader())