diff --git a/src/lazy_loader/__init__.py b/src/lazy_loader/__init__.py index 23b5382..90eddd4 100644 --- a/src/lazy_loader/__init__.py +++ b/src/lazy_loader/__init__.py @@ -5,20 +5,18 @@ Makes it easy to load subpackages and functions on demand. """ -import ast +import _thread import importlib -import importlib.util import os import sys -import threading import types -import warnings __version__ = "0.6rc0.dev0" __all__ = ["attach", "attach_stub", "load"] -threadlock = threading.Lock() +# Same lock type as threading.Lock(), without the threading import cost +threadlock = _thread.allocate_lock() def attach(package_name, submodules=None, submod_attrs=None): @@ -72,23 +70,23 @@ def attach(package_name, submodules=None, submod_attrs=None): def __getattr__(name): if name in submodules: - return importlib.import_module(f"{package_name}.{name}") + attr = 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}") + # Cache the resolved value on the package so that subsequent + # accesses bypass __getattr__; this also ensures an attribute + # shadows a same-named submodule. + pkg = sys.modules.get(package_name) + if pkg is not None: + pkg.__dict__[name] = attr + + return attr + def __dir__(): return __all__.copy() @@ -191,7 +189,11 @@ def myfunc(): if have_module and require is None: return module + import importlib.util + if not suppress_warning and "." in fullname: + import warnings + msg = ( "subpackages can technically be lazily loaded, but it causes the " "package to be eagerly loaded even if it is already lazily loaded. " @@ -276,31 +278,6 @@ def _check_requirement(require: str) -> bool: ) -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. @@ -327,6 +304,32 @@ def attach_stub(package_name: str, filename: str): 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) """ + import ast + + 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) + stubfile = ( filename if filename.endswith("i") else f"{os.path.splitext(filename)[0]}.pyi" ) diff --git a/tests/test_lazy_loader.py b/tests/test_lazy_loader.py index d68537f..be7e4de 100644 --- a/tests/test_lazy_loader.py +++ b/tests/test_lazy_loader.py @@ -178,6 +178,17 @@ def test_attach_same_module_and_attr_name(clean_fake_pkg, eager_import): assert isinstance(some_func, types.FunctionType) +def test_attach_caches_resolved_attrs(clean_fake_pkg): + from tests import fake_pkg + + assert "aux_func" not in vars(fake_pkg) + aux_func = fake_pkg.aux_func + # The resolved attribute is cached on the package, so later accesses + # do not go through __getattr__ again + assert vars(fake_pkg)["aux_func"] is aux_func + assert fake_pkg.aux_func is aux_func + + FAKE_STUB = """ from . import rank from ._gaussian import gaussian