Lazy importing

lazy_import.py

from __future__ import print_function

import re
import imp
import sys
import time
import types
import inspect
import logging
import threading
import importlib
import traceback
import functools
import collections

# pip install six
from six.moves import reload_module

# pip install lazy_object_proxy
from lazy_object_proxy import Proxy


class LazyImportLoader(object):
    # https://peps.python.org/pep-0302/

    _lock = threading.RLock()

    def __init__(self):
        self._names = collections.defaultdict(set)
        self._meta = {}
        self._reg = None
        self._debug = False

    def set_debug(self, state):
        self._debug = state

    def add_module(self, fullname, *attributes):
        self._names[fullname].update(attributes)

    def _build_reg(self):
        self._reg = re.compile(
            r"(?:{})(?:\.[\w\.]+)?$".format("|".join(map(re.escape, self._names)))
        )

    def start(self):
        if self not in sys.meta_path:
            sys.meta_path.insert(0, self)
        return self

    def stop(self):
        while self in sys.meta_path:
            sys.meta_path.remove(self)

    def __enter__(self):
        return self.start()

    def __exit__(self, *_err):
        self.stop()

    def find_module(self, fullname, path=None):
        if fullname in self._meta:
            # Once we have seen the module for the first time
            # we don't need to import again.
            return None

        if not self._reg:
            self._build_reg()

        # Check if we care about this import. Only act if we do.
        if not self._reg.match(fullname):
            return None

        # Track the first time we were asked for the module.
        meta = self._meta[fullname] = {"time": time.time(), "loaded": False}

        # Eagerly search for the module to detect up front it if does
        # not exist. This is so try/except calls can select based
        # on modules existance.
        # We can use some of the package information too.
        package = None
        for name in fullname.split("."):
            result = imp.find_module(name, package)
            if result[0]:  # Close an open file, since we are not using it now
                result[0].close()
            package = result[1] and [result[1]]

        meta["path"] = package
        meta["package"] = fullname.rpartition(".")[0] if bool(result[0]) else fullname
        return self

    def load_module(self, fullname):

        # Build our module proxy
        mod = sys.modules[fullname] = ProxyModule(fullname)
        meta = self._meta[fullname]
        mod.__package__ = meta["package"]
        mod.__path__ = meta["path"]
        mod.__loader__ = self
        for attr in self._names[fullname]:
            setattr(
                mod, attr, Proxy(functools.partial(self.materialize_attr, mod, attr))
            )
        return mod

    def materialize_attr(self, proxy, attribute):
        frame = inspect.currentframe().f_back
        self.materialize_module(mod, frame)
        return getattr(mod, attr)

    def materialize_module(self, proxy, frame):
        with self._lock:
            meta = self._meta[proxy.__name__]
            if meta["loaded"]:
                return

            start = meta["time"]
            current = time.time()

            # Now perform the real import
            reload_module(proxy)
            meta["loaded"] = True

            if self._debug:
                print("".join(traceback.format_stack(frame)), end="")
            print(
                "Imported {} after {:.3f}s delay. Took {:.3f}s.".format(
                    proxy.__name__, current - start, time.time() - current
                )
            )


class ProxyModule(types.ModuleType):
    def __getattribute__(self, name):
        dct = object.__getattribute__(self, "__dict__")
        if name == "__dict__":
            return dct
        if name not in dct:
            frame = inspect.currentframe()
            dct["__loader__"].materialize_module(self, frame.f_back)
        return super(ProxyModule, self).__getattribute__(name)

    def __dir__(self):
        frame = inspect.currentframe()
        self.__loader__.materialize_module(self, frame.f_back)
        return object.__dir__(self)


if __name__ == "__main__":
    with LazyImportLoader() as lazy:
        lazy.add_module("glob", "escape")
        lazy.add_module("multiprocessing.managers")
        lazy.set_debug(True)

        import glob
        from glob import escape

        assert isinstance(glob, ProxyModule)
        assert hasattr(glob, "escape")  # Proxy attribute
        assert "glob" not in glob.__dict__
        assert hasattr(glob, "glob")  # Lazy loading

        import multiprocessing.managers

        assert "multiprocessing" in sys.modules
        assert not isinstance(sys.modules["multiprocessing"], ProxyModule)
        assert isinstance(multiprocessing.managers, ProxyModule)

        from multiprocessing.managers import SyncManager  # Imported...

        assert SyncManager

        try:
            import nothing
        except ImportError:
            pass  # We want to eagerly error, but not eagerly load
        else:
            assert False, "Should have errored"