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
« 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
9from ._hooks import LazyObject
10from ._transformer import TransformModuleImports
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
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)
32enabled_packages = set()
35def scan_distributions():
36 global enabled_packages
37 for dist in importlib.metadata.distributions():
38 metadata = dist.metadata
40 if metadata is None:
41 continue # pragma: no cover
43 if metadata["Keywords"] is None:
44 continue
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)
52def _top_level_declared(dist):
53 return (dist.read_text("top_level.txt") or "").split()
56def _top_level_inferred(dist):
57 files = dist.files
58 if files is None:
59 return {} # pragma: no cover
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 }
67 is_namespace = min(len(p) for p in parts) == 2
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}
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
81 if "LAZY_IMPORTS_LITE_DISABLE" in os.environ:
82 return None
84 spec = super().find_spec(fullname, path, target)
86 if spec is None:
87 return None
89 if spec.origin is None:
90 return None # pragma: no cover
92 name = spec.name.split(".")[0]
93 namespace_name = ".".join(spec.name.split(".")[:2])
95 if (
96 name in enabled_packages or namespace_name in enabled_packages
97 ) and spec.origin.endswith(".py"):
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
110 return None
112 def create_module(self, spec):
113 return LazyModule(spec.name)
115 def exec_module(self, module):
116 origin: str = module.__spec__.origin
118 mod_ast = module.__spec__.mod_ast
119 del module.__spec__.mod_ast
121 transformer = TransformModuleImports()
122 new_ast = transformer.visit(mod_ast)
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"]
131def setup():
132 scan_distributions()
134 if not any(isinstance(m, LazyLoader) for m in sys.meta_path):
135 sys.meta_path.insert(0, LazyLoader())