350 lines
11 KiB
Python
350 lines
11 KiB
Python
|
"""
|
||
|
|
lazy_loader
|
||
|
|
===========
|
||
|
|
|
||
|
|
Makes it easy to load subpackages and functions on demand.
|
||
|
|
"""
|
||
|
|
|
||
|
|
import ast
|
||
|
|
import importlib
|
||
|
|
import importlib.util
|
||
|
|
import os
|
||
|
|
import sys
|
||
|
|
import threading
|
||
|
|
import types
|
||
|
|
import warnings
|
||
|
|
|
||
|
|
__version__ = "0.4"
|
||
|
|
__all__ = ["attach", "load", "attach_stub"]
|
||
|
|
|
||
|
|
|
||
|
|
threadlock = threading.Lock()
|
||
|
|
|
||
|
|
|
||
|
|
def attach(package_name, submodules=None, submod_attrs=None):
|
||
|
|
"""Attach lazily loaded submodules, functions, or other attributes.
|
||
|
|
|
||
|
|
Typically, modules import submodules and attributes as follows::
|
||
|
|
|
||
|
|
import mysubmodule
|
||
|
|
import anothersubmodule
|
||
|
|
|
||
|
|
from .foo import someattr
|
||
|
|
|
||
|
|
The idea is to replace a package's `__getattr__`, `__dir__`, and
|
||
|
|
`__all__`, such that all imports work exactly the way they would
|
||
|
|
with normal imports, except that the import occurs upon first use.
|
||
|
|
|
||
|
|
The typical way to call this function, replacing the above imports, is::
|
||
|
|
|
||
|
|
__getattr__, __dir__, __all__ = lazy.attach(
|
||
|
|
__name__,
|
||
|
|
['mysubmodule', 'anothersubmodule'],
|
||
|
|
{'foo': ['someattr']}
|
||
|
|
)
|
||
|
|
|
||
|
|
This functionality requires Python 3.7 or higher.
|
||
|
|
|
||
|
|
Parameters
|
||
|
|
----------
|
||
|
|
package_name : str
|
||
|
|
Typically use ``__name__``.
|
||
|
|
submodules : set
|
||
|
|
List of submodules to attach.
|
||
|
|
submod_attrs : dict
|
||
|
|
Dictionary of submodule -> list of attributes / functions.
|
||
|
|
These attributes are imported as they are used.
|
||
|
|
|
||
|
|
Returns
|
||
|
|
-------
|
||
|
|
__getattr__, __dir__, __all__
|
||
|
|
|
||
|
|
"""
|
||
|
|
if submod_attrs is None:
|
||
|
|
submod_attrs = {}
|
||
|
|
|
||
|
|
if submodules is None:
|
||
|
|
submodules = set()
|
||
|
|
else:
|
||
|
|
submodules = set(submodules)
|
||
|
|
|
||
|
|
attr_to_modules = {
|
||
|
|
attr: mod for mod, attrs in submod_attrs.items() for attr in attrs
|
||
|
|
}
|
||
|
|
|
||
|
|
__all__ = sorted(submodules | attr_to_modules.keys())
|
||
|
|
|
||
|
|
def __getattr__(name):
|
||
|
|
if name in submodules:
|
||
|
|
return importlib.import_module(f"{package_name}.{name}")
|
||
|
|
elif name in attr_to_modules:
|
||
|
|
submod_path = f"{package_name}.{attr_to_modules[name]}"
|
||
|
|
submod = importlib.import_module(submod_path)
|
||
|
|
attr = getattr(submod, name)
|
||
|
|
|
||
|
|
# If the attribute lives in a file (module) with the same
|
||
|
|
# name as the attribute, ensure that the attribute and *not*
|
||
|
|
# the module is accessible on the package.
|
||
|
|
if name == attr_to_modules[name]:
|
||
|
|
pkg = sys.modules[package_name]
|
||
|
|
pkg.__dict__[name] = attr
|
||
|
|
|
||
|
|
return attr
|
||
|
|
else:
|
||
|
|
raise AttributeError(f"No {package_name} attribute {name}")
|
||
|
|
|
||
|
|
def __dir__():
|
||
|
|
return __all__
|
||
|
|
|
||
|
|
if os.environ.get("EAGER_IMPORT", ""):
|
||
|
|
for attr in set(attr_to_modules.keys()) | submodules:
|
||
|
|
__getattr__(attr)
|
||
|
|
|
||
|
|
return __getattr__, __dir__, list(__all__)
|
||
|
|
|
||
|
|
|
||
|
|
class DelayedImportErrorModule(types.ModuleType):
|
||
|
|
def __init__(self, frame_data, *args, message, **kwargs):
|
||
|
|
self.__frame_data = frame_data
|
||
|
|
self.__message = message
|
||
|
|
super().__init__(*args, **kwargs)
|
||
|
|
|
||
|
|
def __getattr__(self, x):
|
||
|
|
if x in ("__class__", "__file__", "__frame_data", "__message"):
|
||
|
|
super().__getattr__(x)
|
||
|
|
else:
|
||
|
|
fd = self.__frame_data
|
||
|
|
raise ModuleNotFoundError(
|
||
|
|
f"{self.__message}\n\n"
|
||
|
|
"This error is lazily reported, having originally occured in\n"
|
||
|
|
f' File {fd["filename"]}, line {fd["lineno"]}, in {fd["function"]}\n\n'
|
||
|
|
f'----> {"".join(fd["code_context"] or "").strip()}'
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def load(fullname, *, require=None, error_on_import=False):
|
||
|
|
"""Return a lazily imported proxy for a module.
|
||
|
|
|
||
|
|
We often see the following pattern::
|
||
|
|
|
||
|
|
def myfunc():
|
||
|
|
import numpy as np
|
||
|
|
np.norm(...)
|
||
|
|
....
|
||
|
|
|
||
|
|
Putting the import inside the function prevents, in this case,
|
||
|
|
`numpy`, from being imported at function definition time.
|
||
|
|
That saves time if `myfunc` ends up not being called.
|
||
|
|
|
||
|
|
This `load` function returns a proxy module that, upon access, imports
|
||
|
|
the actual module. So the idiom equivalent to the above example is::
|
||
|
|
|
||
|
|
np = lazy.load("numpy")
|
||
|
|
|
||
|
|
def myfunc():
|
||
|
|
np.norm(...)
|
||
|
|
....
|
||
|
|
|
||
|
|
The initial import time is fast because the actual import is delayed
|
||
|
|
until the first attribute is requested. The overall import time may
|
||
|
|
decrease as well for users that don't make use of large portions
|
||
|
|
of your library.
|
||
|
|
|
||
|
|
Warning
|
||
|
|
-------
|
||
|
|
While lazily loading *sub*packages technically works, it causes the
|
||
|
|
package (that contains the subpackage) to be eagerly loaded even
|
||
|
|
if the package is already lazily loaded.
|
||
|
|
So, you probably shouldn't use subpackages with this `load` feature.
|
||
|
|
Instead you should encourage the package maintainers to use the
|
||
|
|
`lazy_loader.attach` to make their subpackages load lazily.
|
||
|
|
|
||
|
|
Parameters
|
||
|
|
----------
|
||
|
|
fullname : str
|
||
|
|
The full name of the module or submodule to import. For example::
|
||
|
|
|
||
|
|
sp = lazy.load('scipy') # import scipy as sp
|
||
|
|
|
||
|
|
require : str
|
||
|
|
A dependency requirement as defined in PEP-508. For example::
|
||
|
|
|
||
|
|
"numpy >=1.24"
|
||
|
|
|
||
|
|
If defined, the proxy module will raise an error if the installed
|
||
|
|
version does not satisfy the requirement.
|
||
|
|
|
||
|
|
error_on_import : bool
|
||
|
|
Whether to postpone raising import errors until the module is accessed.
|
||
|
|
If set to `True`, import errors are raised as soon as `load` is called.
|
||
|
|
|
||
|
|
Returns
|
||
|
|
-------
|
||
|
|
pm : importlib.util._LazyModule
|
||
|
|
Proxy module. Can be used like any regularly imported module.
|
||
|
|
Actual loading of the module occurs upon first attribute request.
|
||
|
|
|
||
|
|
"""
|
||
|
|
with threadlock:
|
||
|
|
module = sys.modules.get(fullname)
|
||
|
|
have_module = module is not None
|
||
|
|
|
||
|
|
# Most common, short-circuit
|
||
|
|
if have_module and require is None:
|
||
|
|
return module
|
||
|
|
|
||
|
|
if "." in fullname:
|
||
|
|
msg = (
|
||
|
|
"subpackages can technically be lazily loaded, but it causes the "
|
||
|
|
"package to be eagerly loaded even if it is already lazily loaded."
|
||
|
|
"So, you probably shouldn't use subpackages with this lazy feature."
|
||
|
|
)
|
||
|
|
warnings.warn(msg, RuntimeWarning)
|
||
|
|
|
||
|
|
spec = None
|
||
|
|
|
||
|
|
if not have_module:
|
||
|
|
spec = importlib.util.find_spec(fullname)
|
||
|
|
have_module = spec is not None
|
||
|
|
|
||
|
|
if not have_module:
|
||
|
|
not_found_message = f"No module named '{fullname}'"
|
||
|
|
elif require is not None:
|
||
|
|
try:
|
||
|
|
have_module = _check_requirement(require)
|
||
|
|
except ModuleNotFoundError as e:
|
||
|
|
raise ValueError(
|
||
|
|
f"Found module '{fullname}' but cannot test "
|
||
|
|
"requirement '{require}'. "
|
||
|
|
"Requirements must match distribution name, not module name."
|
||
|
|
) from e
|
||
|
|
|
||
|
|
not_found_message = f"No distribution can be found matching '{require}'"
|
||
|
|
|
||
|
|
if not have_module:
|
||
|
|
if error_on_import:
|
||
|
|
raise ModuleNotFoundError(not_found_message)
|
||
|
|
import inspect
|
||
|
|
|
||
|
|
try:
|
||
|
|
parent = inspect.stack()[1]
|
||
|
|
frame_data = {
|
||
|
|
"filename": parent.filename,
|
||
|
|
"lineno": parent.lineno,
|
||
|
|
"function": parent.function,
|
||
|
|
"code_context": parent.code_context,
|
||
|
|
}
|
||
|
|
return DelayedImportErrorModule(
|
||
|
|
frame_data,
|
||
|
|
"DelayedImportErrorModule",
|
||
|
|
message=not_found_message,
|
||
|
|
)
|
||
|
|
finally:
|
||
|
|
del parent
|
||
|
|
|
||
|
|
if spec is not None:
|
||
|
|
module = importlib.util.module_from_spec(spec)
|
||
|
|
sys.modules[fullname] = module
|
||
|
|
|
||
|
|
loader = importlib.util.LazyLoader(spec.loader)
|
||
|
|
loader.exec_module(module)
|
||
|
|
|
||
|
|
return module
|
||
|
|
|
||
|
|
|
||
|
|
def _check_requirement(require: str) -> bool:
|
||
|
|
"""Verify that a package requirement is satisfied
|
||
|
|
|
||
|
|
If the package is required, a ``ModuleNotFoundError`` is raised
|
||
|
|
by ``importlib.metadata``.
|
||
|
|
|
||
|
|
Parameters
|
||
|
|
----------
|
||
|
|
require : str
|
||
|
|
A dependency requirement as defined in PEP-508
|
||
|
|
|
||
|
|
Returns
|
||
|
|
-------
|
||
|
|
satisfied : bool
|
||
|
|
True if the installed version of the dependency matches
|
||
|
|
the specified version, False otherwise.
|
||
|
|
"""
|
||
|
|
import packaging.requirements
|
||
|
|
|
||
|
|
try:
|
||
|
|
import importlib.metadata as importlib_metadata
|
||
|
|
except ImportError: # PY37
|
||
|
|
import importlib_metadata
|
||
|
|
|
||
|
|
req = packaging.requirements.Requirement(require)
|
||
|
|
return req.specifier.contains(
|
||
|
|
importlib_metadata.version(req.name),
|
||
|
|
prereleases=True,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
class _StubVisitor(ast.NodeVisitor):
|
||
|
|
"""AST visitor to parse a stub file for submodules and submod_attrs."""
|
||
|
|
|
||
|
|
def __init__(self):
|
||
|
|
self._submodules = set()
|
||
|
|
self._submod_attrs = {}
|
||
|
|
|
||
|
|
def visit_ImportFrom(self, node: ast.ImportFrom):
|
||
|
|
if node.level != 1:
|
||
|
|
raise ValueError(
|
||
|
|
"Only within-module imports are supported (`from .* import`)"
|
||
|
|
)
|
||
|
|
if node.module:
|
||
|
|
attrs: list = self._submod_attrs.setdefault(node.module, [])
|
||
|
|
aliases = [alias.name for alias in node.names]
|
||
|
|
if "*" in aliases:
|
||
|
|
raise ValueError(
|
||
|
|
"lazy stub loader does not support star import "
|
||
|
|
f"`from {node.module} import *`"
|
||
|
|
)
|
||
|
|
attrs.extend(aliases)
|
||
|
|
else:
|
||
|
|
self._submodules.update(alias.name for alias in node.names)
|
||
|
|
|
||
|
|
|
||
|
|
def attach_stub(package_name: str, filename: str):
|
||
|
|
"""Attach lazily loaded submodules, functions from a type stub.
|
||
|
|
|
||
|
|
This is a variant on ``attach`` that will parse a `.pyi` stub file to
|
||
|
|
infer ``submodules`` and ``submod_attrs``. This allows static type checkers
|
||
|
|
to find imports, while still providing lazy loading at runtime.
|
||
|
|
|
||
|
|
Parameters
|
||
|
|
----------
|
||
|
|
package_name : str
|
||
|
|
Typically use ``__name__``.
|
||
|
|
filename : str
|
||
|
|
Path to `.py` file which has an adjacent `.pyi` file.
|
||
|
|
Typically use ``__file__``.
|
||
|
|
|
||
|
|
Returns
|
||
|
|
-------
|
||
|
|
__getattr__, __dir__, __all__
|
||
|
|
The same output as ``attach``.
|
||
|
|
|
||
|
|
Raises
|
||
|
|
------
|
||
|
|
ValueError
|
||
|
|
If a stub file is not found for `filename`, or if the stubfile is formmated
|
||
|
|
incorrectly (e.g. if it contains an relative import from outside of the module)
|
||
|
|
"""
|
||
|
|
stubfile = (
|
||
|
|
filename if filename.endswith("i") else f"{os.path.splitext(filename)[0]}.pyi"
|
||
|
|
)
|
||
|
|
|
||
|
|
if not os.path.exists(stubfile):
|
||
|
|
raise ValueError(f"Cannot load imports from non-existent stub {stubfile!r}")
|
||
|
|
|
||
|
|
with open(stubfile) as f:
|
||
|
|
stub_node = ast.parse(f.read())
|
||
|
|
|
||
|
|
visitor = _StubVisitor()
|
||
|
|
visitor.visit(stub_node)
|
||
|
|
return attach(package_name, visitor._submodules, visitor._submod_attrs)
|