Lazy Dispatch Test

%load_ext autoreload
%autoreload 2

class A:
    class A:
        """A test class."""


class B:
    class B:
        """A test class."""


class C(A):
    class C(A.A):
        """A test class."""


class D(A, B):
    pass


class E(C, D):
    pass

import torch

torch.nn.Module

torch.Tensor

from lazy_dispatch.isinstance import lazy_isinstance

t = torch.tensor([1, 2, 3])
lazy_isinstance(t, "torch.Tensor")

True

from typing import TYPE_CHECKING

import lazy_dispatch as ld

if TYPE_CHECKING:
    import keras  # type: ignore  # noqa: PGH003
    import torch


@ld.lazy_singledispatch
def f(_: object) -> str:
    return "unknown"


@f.register
def _(_: int) -> str:
    return "int"


@f.register
def _(_: "torch.nn.modules.module.Module") -> str:
    return "torch module"


@f.register
def _(_: "torch.nn.modules.linear.Linear") -> str:
    return "torch linear"


@f.register
def _(_: "keras.engine.training.Model") -> str:
    return "keras model"


@f.register
def _(_: A) -> str:
    return "A"


@f.register
def _(_: "__main__.B") -> str:  # type: ignore  # noqa: F821, PGH003
    return "B"


@f.register
def _(_: E) -> str:
    return "E"

f.__annotations__

{'_': object, 'return': str}

import torch.nn

layer = torch.nn.Linear(10, 10)
seq = torch.nn.Sequential(layer, layer)

print(42, f(42))
print(1.5, f(1.5))
print("linear:", f(layer))
print("seq:", f(seq))
print("A:", f(A()))
print("B:", f(B()))
print("C:", f(C()))
print("D:", f(D()))
print("E:", f(E()))

42 int
1.5 unknown
linear: torch linear
seq: torch module
A: A
B: B
C: A
D: A
E: E

f.string_registry

mappingproxy({'keras.engine.training.Model': <function __main__._(_: 'keras.engine.training.Model') -> str>})