lazy_imports.py 5.6 KB

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