Question

How to make a lazy loader play nice with static type checking?

I've written a crude lazy importer so you can do stuff like this:

from loader import Lazy

httpx = Lazy("httpx")  # The `httpx` module is not yet loaded

httpx.get("https://google.ca/")  # Loads the `httpx` module and calls `.get()`

This is what it looks like:

from functools import cached_property
from importlib import import_module
from typing import Any, TYPE_CHECKING


class Lazy:
    def __init__(self, name: str) -> None:
        self._name = name

    def __getattr__(self, item: str) -> Any:
        return getattr(self._module, item)

    @cached_property
    def _module(self):
        return import_module(self._name)

...and it works! However, static type checkers like Mypy and PyCharm treat the lazily-imported httpx module as if it's capable of anything, so code like this:

from loader import Lazy

httpx = Lazy("httpx")
httpx.get(42)
httpx.woot

...isn't flagged for being broken. PyCharm has no way of autocompleting method names or arguments either, so while the code runs, it's much harder to develop on.

In a perfect world, the lazy loader would have a way of swapping itself out for the lazily-imported module in the eyes of the static type checker, but that'd require the static typechecker to do dynamic things, so I'm not even sure that can be done.

Is there an option available to me, or is this simply a no-no in Pythonland? Normally, I wouldn't even try to do something like this, but the codebase I'm working on is Very Big and a lazy loader could got a long way to improving start times when doing local development.

 4  90  4
1 Jan 1970

Solution

 1

@InSync has provided a good explanation of why this is not very simple. However, you mention lazy loading would be just for local development. Therefore, you may be okay with taking on the below higher-risk solution. Here is a good start on how you can actually translate some import statements into their lazy equivalent. Because the import statements stay as-is, mypy / PyCharm type checking works just as you'd expect.

example.py

print("Loaded.")

def call() -> None:
    print("Called.")

main.py (imports in playground below)

class LazyModule(ModuleType):
    def __getattr__(self, item: str) -> Any:
        return getattr(self._module, item)

    @cached_property
    def _module(self) -> ModuleType:
        importlib.reload(self)
        return import_module(self.__name__)

# Ported from https://github.com/python/cpython/blob/main/Lib/importlib/_bootstrap.py.
def init_module_attributes(module: ModuleType, spec: ModuleSpec) -> None:
    module.__name__ = spec.name
    module.__loader__ = LAZY_LOADER
    module.__package__ = spec.parent
    module.__spec__ = spec
    module.__path__ = (spec.submodule_search_locations or [])[:]

class LazyLoader(Loader):
    def create_module(self, spec: ModuleSpec) -> ModuleType:
        module = LazyModule(spec.name)

        init_module_attributes(module, spec)
        return module

    def exec_module(self, module: ModuleType) -> None:
        pass

class LazyModuleFinder(MetaPathFinder):
    def find_spec(self, name: str, path: Sequence[str] | None, target: ModuleType | None = None) -> ModuleSpec | None:
        return importlib.util.spec_from_loader(name, LAZY_LOADER)

LAZY_LOADER: Final = LazyLoader()
LAZY_MODULE_FINDER: Final = LazyModuleFinder()

@contextmanager
def lazy_init() -> Iterator[None]:
    try:
        sys.meta_path.insert(0, LAZY_MODULE_FINDER)
        yield
    finally:
        sys.meta_path.remove(LAZY_MODULE_FINDER)

with lazy_init():
    import example

print("Imported.")
example.call()

If you're interested in this approach, I'd first like you to test importing your own modules under lazy_init before I invest into editing the answer and explaining everything. For the simple example module I created, I get the below output when running the above script:

Imported.
Loaded.
Called.

showing that all module initialization is successfully deferred to first attribute access.

This solution itself also passes mypy type checking (with the one error being missing example since I can't attach that in playground), and is designed to be as drop-in as possible. You only need to import your modules under lazy_init.

2024-07-06
Mario Ishac

Solution

 0

All modules are of type types.ModuleType, but each module has its own attributes. You can type hint a function as returning a value of this type, but you can't (yet) use type hints to represent a specific instance, aside from those allowed as arguments to Literal[].

As an analogy, that would be similar to:

class C:
    pass

a = C()
a.foo = 'bar'
def f() -> C:  # Literal[a] is invalid.
    return a

c = f()
c.foo  # error

As long as you maintain consistency with the actual module, you can lie by importing directly only at type-checking time:

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    import httpx
else:
    httpx = Lazy('httpx')

If there are derivations (the _module attribute, for example), handle them with cast() and/or type: ignore. Do note that PyCharm will not respect this guard and will try to analyze both branches. This is hardly a downside, however.

If you have a lot of derivations and only use Mypy, you can also write a Mypy plugin. A Mypy plugin can, theoretically, support any kinds of dynamic behaviours. The downside is that it won't work with other type checkers, and extra maintenance efforts are needed.

2024-07-05
InSync