lazy_imports.py 5.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188
  1. import importlib
  2. import importlib.util
  3. import inspect
  4. import os
  5. import sys
  6. import types
  7. __all__ = ["attach", "_lazy_import"]
  8. def attach(module_name, submodules=None, submod_attrs=None):
  9. """Attach lazily loaded submodules, and functions or other attributes.
  10. Typically, modules import submodules and attributes as follows::
  11. import mysubmodule
  12. import anothersubmodule
  13. from .foo import someattr
  14. The idea of this function is to replace the `__init__.py`
  15. module's `__getattr__`, `__dir__`, and `__all__` attributes such that
  16. all imports work exactly the way they normally would, except that the
  17. actual import is delayed until the resulting module object is first used.
  18. The typical way to call this function, replacing the above imports, is::
  19. __getattr__, __lazy_dir__, __all__ = lazy.attach(
  20. __name__, ["mysubmodule", "anothersubmodule"], {"foo": "someattr"}
  21. )
  22. This functionality requires Python 3.7 or higher.
  23. Parameters
  24. ----------
  25. module_name : str
  26. Typically use __name__.
  27. submodules : set
  28. List of submodules to lazily import.
  29. submod_attrs : dict
  30. Dictionary of submodule -> list of attributes / functions.
  31. These attributes are imported as they are used.
  32. Returns
  33. -------
  34. __getattr__, __dir__, __all__
  35. """
  36. if submod_attrs is None:
  37. submod_attrs = {}
  38. if submodules is None:
  39. submodules = set()
  40. else:
  41. submodules = set(submodules)
  42. attr_to_modules = {
  43. attr: mod for mod, attrs in submod_attrs.items() for attr in attrs
  44. }
  45. __all__ = list(submodules | attr_to_modules.keys())
  46. def __getattr__(name):
  47. if name in submodules:
  48. return importlib.import_module(f"{module_name}.{name}")
  49. elif name in attr_to_modules:
  50. submod = importlib.import_module(f"{module_name}.{attr_to_modules[name]}")
  51. return getattr(submod, name)
  52. else:
  53. raise AttributeError(f"No {module_name} attribute {name}")
  54. def __dir__():
  55. return __all__
  56. if os.environ.get("EAGER_IMPORT", ""):
  57. for attr in set(attr_to_modules.keys()) | submodules:
  58. __getattr__(attr)
  59. return __getattr__, __dir__, list(__all__)
  60. class DelayedImportErrorModule(types.ModuleType):
  61. def __init__(self, frame_data, *args, **kwargs):
  62. self.__frame_data = frame_data
  63. super().__init__(*args, **kwargs)
  64. def __getattr__(self, x):
  65. if x in ("__class__", "__file__", "__frame_data"):
  66. super().__getattr__(x)
  67. else:
  68. fd = self.__frame_data
  69. raise ModuleNotFoundError(
  70. f"No module named '{fd['spec']}'\n\n"
  71. "This error is lazily reported, having originally occurred in\n"
  72. f" File {fd['filename']}, line {fd['lineno']}, in {fd['function']}\n\n"
  73. f"----> {''.join(fd['code_context'] or '').strip()}"
  74. )
  75. def _lazy_import(fullname):
  76. """Return a lazily imported proxy for a module or library.
  77. Warning
  78. -------
  79. Importing using this function can currently cause trouble
  80. when the user tries to import from a subpackage of a module before
  81. the package is fully imported. In particular, this idiom may not work:
  82. np = lazy_import("numpy")
  83. from numpy.lib import recfunctions
  84. This is due to a difference in the way Python's LazyLoader handles
  85. subpackage imports compared to the normal import process. Hopefully
  86. we will get Python's LazyLoader to fix this, or find a workaround.
  87. In the meantime, this is a potential problem.
  88. The workaround is to import numpy before importing from the subpackage.
  89. Notes
  90. -----
  91. We often see the following pattern::
  92. def myfunc():
  93. import scipy as sp
  94. sp.argmin(...)
  95. ....
  96. This is to prevent a library, in this case `scipy`, from being
  97. imported at function definition time, since that can be slow.
  98. This function provides a proxy module that, upon access, imports
  99. the actual module. So the idiom equivalent to the above example is::
  100. sp = lazy.load("scipy")
  101. def myfunc():
  102. sp.argmin(...)
  103. ....
  104. The initial import time is fast because the actual import is delayed
  105. until the first attribute is requested. The overall import time may
  106. decrease as well for users that don't make use of large portions
  107. of the library.
  108. Parameters
  109. ----------
  110. fullname : str
  111. The full name of the package or subpackage to import. For example::
  112. sp = lazy.load("scipy") # import scipy as sp
  113. spla = lazy.load("scipy.linalg") # import scipy.linalg as spla
  114. Returns
  115. -------
  116. pm : importlib.util._LazyModule
  117. Proxy module. Can be used like any regularly imported module.
  118. Actual loading of the module occurs upon first attribute request.
  119. """
  120. try:
  121. return sys.modules[fullname]
  122. except:
  123. pass
  124. # Not previously loaded -- look it up
  125. spec = importlib.util.find_spec(fullname)
  126. if spec is None:
  127. try:
  128. parent = inspect.stack()[1]
  129. frame_data = {
  130. "spec": fullname,
  131. "filename": parent.filename,
  132. "lineno": parent.lineno,
  133. "function": parent.function,
  134. "code_context": parent.code_context,
  135. }
  136. return DelayedImportErrorModule(frame_data, "DelayedImportErrorModule")
  137. finally:
  138. del parent
  139. module = importlib.util.module_from_spec(spec)
  140. sys.modules[fullname] = module
  141. loader = importlib.util.LazyLoader(spec.loader)
  142. loader.exec_module(module)
  143. return module