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

73 statements  

« prev     ^ index     » next       coverage.py v7.4.1, created at 2024-02-11 15:33 +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 return { 

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

63 for f in files 

64 if f.suffix == ".py" 

65 } 

66 

67 

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

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

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

71 # fix wired windows bug 

72 return None 

73 

74 if "LAZY_IMPORTS_LITE_DISABLE" in os.environ: 

75 return None 

76 

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

78 

79 if spec is None: 

80 return None 

81 

82 if spec.origin is None: 

83 return None # pragma: no cover 

84 

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

86 

87 if name in enabled_packages and spec.origin.endswith(".py"): 

88 spec.loader = self 

89 return spec 

90 

91 return None 

92 

93 def create_module(self, spec): 

94 return LazyModule(spec.name) 

95 

96 def exec_module(self, module): 

97 origin: str = module.__spec__.origin 

98 with open(origin) as f: 

99 mod_raw = f.read() 

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

101 transformer = TransformModuleImports() 

102 new_ast = transformer.visit(mod_ast) 

103 

104 ast.fix_missing_locations(new_ast) 

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

106 exec(mod_code, module.__dict__) 

107 del module.__dict__["__lazy_imports_lite__"] 

108 del module.__dict__["globals"] 

109 

110 

111def setup(): 

112 scan_distributions() 

113 

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

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