__init__.py 11 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349
  1. """
  2. lazy_loader
  3. ===========
  4. Makes it easy to load subpackages and functions on demand.
  5. """
  6. import ast
  7. import importlib
  8. import importlib.util
  9. import os
  10. import sys
  11. import threading
  12. import types
  13. import warnings
  14. __version__ = "0.4"
  15. __all__ = ["attach", "load", "attach_stub"]
  16. threadlock = threading.Lock()
  17. def attach(package_name, submodules=None, submod_attrs=None):
  18. """Attach lazily loaded submodules, functions, or other attributes.
  19. Typically, modules import submodules and attributes as follows::
  20. import mysubmodule
  21. import anothersubmodule
  22. from .foo import someattr
  23. The idea is to replace a package's `__getattr__`, `__dir__`, and
  24. `__all__`, such that all imports work exactly the way they would
  25. with normal imports, except that the import occurs upon first use.
  26. The typical way to call this function, replacing the above imports, is::
  27. __getattr__, __dir__, __all__ = lazy.attach(
  28. __name__,
  29. ['mysubmodule', 'anothersubmodule'],
  30. {'foo': ['someattr']}
  31. )
  32. This functionality requires Python 3.7 or higher.
  33. Parameters
  34. ----------
  35. package_name : str
  36. Typically use ``__name__``.
  37. submodules : set
  38. List of submodules to attach.
  39. submod_attrs : dict
  40. Dictionary of submodule -> list of attributes / functions.
  41. These attributes are imported as they are used.
  42. Returns
  43. -------
  44. __getattr__, __dir__, __all__
  45. """
  46. if submod_attrs is None:
  47. submod_attrs = {}
  48. if submodules is None:
  49. submodules = set()
  50. else:
  51. submodules = set(submodules)
  52. attr_to_modules = {
  53. attr: mod for mod, attrs in submod_attrs.items() for attr in attrs
  54. }
  55. __all__ = sorted(submodules | attr_to_modules.keys())
  56. def __getattr__(name):
  57. if name in submodules:
  58. return importlib.import_module(f"{package_name}.{name}")
  59. elif name in attr_to_modules:
  60. submod_path = f"{package_name}.{attr_to_modules[name]}"
  61. submod = importlib.import_module(submod_path)
  62. attr = getattr(submod, name)
  63. # If the attribute lives in a file (module) with the same
  64. # name as the attribute, ensure that the attribute and *not*
  65. # the module is accessible on the package.
  66. if name == attr_to_modules[name]:
  67. pkg = sys.modules[package_name]
  68. pkg.__dict__[name] = attr
  69. return attr
  70. else:
  71. raise AttributeError(f"No {package_name} attribute {name}")
  72. def __dir__():
  73. return __all__
  74. if os.environ.get("EAGER_IMPORT", ""):
  75. for attr in set(attr_to_modules.keys()) | submodules:
  76. __getattr__(attr)
  77. return __getattr__, __dir__, list(__all__)
  78. class DelayedImportErrorModule(types.ModuleType):
  79. def __init__(self, frame_data, *args, message, **kwargs):
  80. self.__frame_data = frame_data
  81. self.__message = message
  82. super().__init__(*args, **kwargs)
  83. def __getattr__(self, x):
  84. if x in ("__class__", "__file__", "__frame_data", "__message"):
  85. super().__getattr__(x)
  86. else:
  87. fd = self.__frame_data
  88. raise ModuleNotFoundError(
  89. f"{self.__message}\n\n"
  90. "This error is lazily reported, having originally occured in\n"
  91. f' File {fd["filename"]}, line {fd["lineno"]}, in {fd["function"]}\n\n'
  92. f'----> {"".join(fd["code_context"] or "").strip()}'
  93. )
  94. def load(fullname, *, require=None, error_on_import=False):
  95. """Return a lazily imported proxy for a module.
  96. We often see the following pattern::
  97. def myfunc():
  98. import numpy as np
  99. np.norm(...)
  100. ....
  101. Putting the import inside the function prevents, in this case,
  102. `numpy`, from being imported at function definition time.
  103. That saves time if `myfunc` ends up not being called.
  104. This `load` function returns a proxy module that, upon access, imports
  105. the actual module. So the idiom equivalent to the above example is::
  106. np = lazy.load("numpy")
  107. def myfunc():
  108. np.norm(...)
  109. ....
  110. The initial import time is fast because the actual import is delayed
  111. until the first attribute is requested. The overall import time may
  112. decrease as well for users that don't make use of large portions
  113. of your library.
  114. Warning
  115. -------
  116. While lazily loading *sub*packages technically works, it causes the
  117. package (that contains the subpackage) to be eagerly loaded even
  118. if the package is already lazily loaded.
  119. So, you probably shouldn't use subpackages with this `load` feature.
  120. Instead you should encourage the package maintainers to use the
  121. `lazy_loader.attach` to make their subpackages load lazily.
  122. Parameters
  123. ----------
  124. fullname : str
  125. The full name of the module or submodule to import. For example::
  126. sp = lazy.load('scipy') # import scipy as sp
  127. require : str
  128. A dependency requirement as defined in PEP-508. For example::
  129. "numpy >=1.24"
  130. If defined, the proxy module will raise an error if the installed
  131. version does not satisfy the requirement.
  132. error_on_import : bool
  133. Whether to postpone raising import errors until the module is accessed.
  134. If set to `True`, import errors are raised as soon as `load` is called.
  135. Returns
  136. -------
  137. pm : importlib.util._LazyModule
  138. Proxy module. Can be used like any regularly imported module.
  139. Actual loading of the module occurs upon first attribute request.
  140. """
  141. with threadlock:
  142. module = sys.modules.get(fullname)
  143. have_module = module is not None
  144. # Most common, short-circuit
  145. if have_module and require is None:
  146. return module
  147. if "." in fullname:
  148. msg = (
  149. "subpackages can technically be lazily loaded, but it causes the "
  150. "package to be eagerly loaded even if it is already lazily loaded."
  151. "So, you probably shouldn't use subpackages with this lazy feature."
  152. )
  153. warnings.warn(msg, RuntimeWarning)
  154. spec = None
  155. if not have_module:
  156. spec = importlib.util.find_spec(fullname)
  157. have_module = spec is not None
  158. if not have_module:
  159. not_found_message = f"No module named '{fullname}'"
  160. elif require is not None:
  161. try:
  162. have_module = _check_requirement(require)
  163. except ModuleNotFoundError as e:
  164. raise ValueError(
  165. f"Found module '{fullname}' but cannot test "
  166. "requirement '{require}'. "
  167. "Requirements must match distribution name, not module name."
  168. ) from e
  169. not_found_message = f"No distribution can be found matching '{require}'"
  170. if not have_module:
  171. if error_on_import:
  172. raise ModuleNotFoundError(not_found_message)
  173. import inspect
  174. try:
  175. parent = inspect.stack()[1]
  176. frame_data = {
  177. "filename": parent.filename,
  178. "lineno": parent.lineno,
  179. "function": parent.function,
  180. "code_context": parent.code_context,
  181. }
  182. return DelayedImportErrorModule(
  183. frame_data,
  184. "DelayedImportErrorModule",
  185. message=not_found_message,
  186. )
  187. finally:
  188. del parent
  189. if spec is not None:
  190. module = importlib.util.module_from_spec(spec)
  191. sys.modules[fullname] = module
  192. loader = importlib.util.LazyLoader(spec.loader)
  193. loader.exec_module(module)
  194. return module
  195. def _check_requirement(require: str) -> bool:
  196. """Verify that a package requirement is satisfied
  197. If the package is required, a ``ModuleNotFoundError`` is raised
  198. by ``importlib.metadata``.
  199. Parameters
  200. ----------
  201. require : str
  202. A dependency requirement as defined in PEP-508
  203. Returns
  204. -------
  205. satisfied : bool
  206. True if the installed version of the dependency matches
  207. the specified version, False otherwise.
  208. """
  209. import packaging.requirements
  210. try:
  211. import importlib.metadata as importlib_metadata
  212. except ImportError: # PY37
  213. import importlib_metadata
  214. req = packaging.requirements.Requirement(require)
  215. return req.specifier.contains(
  216. importlib_metadata.version(req.name),
  217. prereleases=True,
  218. )
  219. class _StubVisitor(ast.NodeVisitor):
  220. """AST visitor to parse a stub file for submodules and submod_attrs."""
  221. def __init__(self):
  222. self._submodules = set()
  223. self._submod_attrs = {}
  224. def visit_ImportFrom(self, node: ast.ImportFrom):
  225. if node.level != 1:
  226. raise ValueError(
  227. "Only within-module imports are supported (`from .* import`)"
  228. )
  229. if node.module:
  230. attrs: list = self._submod_attrs.setdefault(node.module, [])
  231. aliases = [alias.name for alias in node.names]
  232. if "*" in aliases:
  233. raise ValueError(
  234. "lazy stub loader does not support star import "
  235. f"`from {node.module} import *`"
  236. )
  237. attrs.extend(aliases)
  238. else:
  239. self._submodules.update(alias.name for alias in node.names)
  240. def attach_stub(package_name: str, filename: str):
  241. """Attach lazily loaded submodules, functions from a type stub.
  242. This is a variant on ``attach`` that will parse a `.pyi` stub file to
  243. infer ``submodules`` and ``submod_attrs``. This allows static type checkers
  244. to find imports, while still providing lazy loading at runtime.
  245. Parameters
  246. ----------
  247. package_name : str
  248. Typically use ``__name__``.
  249. filename : str
  250. Path to `.py` file which has an adjacent `.pyi` file.
  251. Typically use ``__file__``.
  252. Returns
  253. -------
  254. __getattr__, __dir__, __all__
  255. The same output as ``attach``.
  256. Raises
  257. ------
  258. ValueError
  259. If a stub file is not found for `filename`, or if the stubfile is formmated
  260. incorrectly (e.g. if it contains an relative import from outside of the module)
  261. """
  262. stubfile = (
  263. filename if filename.endswith("i") else f"{os.path.splitext(filename)[0]}.pyi"
  264. )
  265. if not os.path.exists(stubfile):
  266. raise ValueError(f"Cannot load imports from non-existent stub {stubfile!r}")
  267. with open(stubfile) as f:
  268. stub_node = ast.parse(f.read())
  269. visitor = _StubVisitor()
  270. visitor.visit(stub_node)
  271. return attach(package_name, visitor._submodules, visitor._submod_attrs)