Possible type-safe LazyObject

[jlevy]

Recently found lazyasd and it seems like a great way to reduce startup times etc. But I found it problematic to use @lazyobject without type warnings.

Below is one approach to fix this. Curious if others find this is a good solution.

If so, this updated, typed version of @lazyobject may be useful.

from typing import Any, Callable, Dict, Generic, Iterator, TypeVar, Mapping, cast

T = TypeVar("T")

class LazyObject(Generic[T]):
    def __init__(self, load: Callable[[], T], ctx: Mapping[str, T], name: str):
        """
        Lazily loads an object via the load function the first time an
        attribute is accessed. Once loaded it will replace itself in the
        provided context (typically the globals of the call site) with the
        given name.

        Parameters
        ----------
        load : Callable[[], T]
            A loader function that performs the actual object construction.
        ctx : Mapping[str, T]
            Context to replace the LazyObject instance in
            with the object returned by load().
        name : str
            Name in the context to give the loaded object. This *should*
            be the name on the LHS of the assignment.
        """
        self._lasdo: Dict[str, Any] = {
            "loaded": False,
            "load": load,
            "ctx": ctx,
            "name": name,
        }

    def _lazy_obj(self) -> T:
        d = self._lasdo
        if d["loaded"]:
            return d["obj"]
        try:
            obj = d["load"]()
            d["ctx"][d["name"]] = d["obj"] = obj
            d["loaded"] = True
            return obj
        except Exception as e:
            raise RuntimeError(f"Error loading object: {e}")

    def __getattribute__(self, name: str) -> Any:
        if name in {"_lasdo", "_lazy_obj"}:
            return super().__getattribute__(name)
        obj = self._lazy_obj()
        return getattr(obj, name)

    def __bool__(self) -> bool:
        return bool(self._lazy_obj())

    def __iter__(self) -> Iterator:
        return iter(self._lazy_obj())  # type: ignore

    def __getitem__(self, item: Any) -> Any:
        return self._lazy_obj()[item]  # type: ignore

    def __setitem__(self, key: Any, value: Any) -> None:
        self._lazy_obj()[key] = value  # type: ignore

    def __delitem__(self, item: Any) -> None:
        del self._lazy_obj()[item]  # type: ignore

    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        return self._lazy_obj()(*args, **kwargs)  # type: ignore

    def __lt__(self, other: Any) -> bool:
        return self._lazy_obj() < other

    def __le__(self, other: Any) -> bool:
        return self._lazy_obj() <= other

    def __eq__(self, other: Any) -> bool:
        return self._lazy_obj() == other

    def __ne__(self, other: Any) -> bool:
        return self._lazy_obj() != other

    def __gt__(self, other: Any) -> bool:
        return self._lazy_obj() > other

    def __ge__(self, other: Any) -> bool:
        return self._lazy_obj() >= other

    def __hash__(self) -> int:
        return hash(self._lazy_obj())

    def __or__(self, other: Any) -> Any:
        return self._lazy_obj() | other

    def __str__(self) -> str:
        return str(self._lazy_obj())

    def __repr__(self) -> str:
        return repr(self._lazy_obj())


def lazyobject(f: Callable[[], T]) -> T:
    """
    Decorator for constructing lazy objects from a function.

    For simplicity, we tell a white lie to the type checker that this is actually of type T.
    """
    return cast(T, LazyObject(f, f.__globals__, f.__name__))