Fixed jaxtyped breaking descriptors. Fixed long module names. (#25)

jaxtyping/__init__.py

@@ -53,4 +53,4 @@ from .import_hook import install_import_hook
 from .pytree_type import PyTree
 
 
-__version__ = "0.2.1"
+__version__ = "0.2.2"

jaxtyping/array_types.py

@@ -388,7 +388,7 @@ class _MetaAbstractDtype(type):
             name = "Array"
         else:
             raise ValueError(f"array_name_format {_array_name_format} not recognised")
-        return _MetaAbstractArray(
+        out = _MetaAbstractArray(
             name,
             (AbstractArray,),
             dict(
@@ -398,6 +398,8 @@ class _MetaAbstractDtype(type):
                 index_variadic=index_variadic,
             ),
         )
+        out.__module__ = "jaxtyping"
+        return out
 
 
 class AbstractDtype(metaclass=_MetaAbstractDtype):
@@ -468,6 +470,7 @@ else:
 
         _Cls.__name__ = name
         _Cls.__qualname__ = name
+        _Cls.__module__ = "jaxtyping"
         return _Cls
 
     UInt8 = _make_dtype(_uint8, "UInt8")

jaxtyping/decorator.py

@@ -24,17 +24,24 @@ import threading
 storage = threading.local()
 
 
-def jaxtyped(fn):
-    @ft.wraps(fn)
-    def wrapper(*args, **kwargs):
+class _Jaxtyped:
+    def __init__(self, fn):
+        self.fn = fn
+
+    def __get__(self, instance, owner):
+        return ft.wraps(self.fn)(_Jaxtyped(self.fn.__get__(instance, owner)))
+
+    def __call__(self, *args, **kwargs):
         try:
             memo_stack = storage.memo_stack
         except AttributeError:
             memo_stack = storage.memo_stack = []
         memo_stack.append(({}, {}, {}))
         try:
-            return fn(*args, **kwargs)
+            return self.fn(*args, **kwargs)
         finally:
             memo_stack.pop()
 
-    return wrapper
+
+def jaxtyped(fn):
+    return ft.wraps(fn)(_Jaxtyped(fn))

jaxtyping/pytree_type.py

@@ -33,6 +33,7 @@ class _FakePyTree(Generic[_T]):
 
 _FakePyTree.__name__ = "PyTree"
 _FakePyTree.__qualname__ = "PyTree"
+_FakePyTree.__module__ = "builtins"
 # Can't do type("PyTree", (Generic[_T],), {}) because dynamic subclassing of typeforms
 # isn't allowed.
 # Can't do types.new_class("PyTree", (Generic[_T],), {}) because that has __module__
@@ -49,7 +50,9 @@ class _MetaPyTree(type):
     @ft.lru_cache(maxsize=None)
     def __getitem__(cls, item):
         name = str(_FakePyTree[item])
-        return _MetaSubscriptPyTree(name, (), {"leaftype": item})
+        out = _MetaSubscriptPyTree(name, (), {"leaftype": item})
+        out.__module__ = "jaxtyping"
+        return out
 
 
 class _MetaSubscriptPyTree(type):
@@ -80,6 +83,7 @@ class _MetaSubscriptPyTree(type):
 
 
 PyTree = _MetaPyTree("PyTree", (), {})
+PyTree.__module__ = "jaxtyping"
 # Can't do `class PyTree(Generic[_T]): ...` because we need to override the
 # instancecheck for PyTree[foo], but we subclassing
 # `type(Generic[int])`, i.e. `typing._GenericAlias` is disallowed.

test/test_decorator.py

@@ -0,0 +1,13 @@
+from jaxtyping import jaxtyped
+
+
+class M:
+    @jaxtyped
+    @classmethod
+    def f(cls):
+        return 3
+
+
+# Check that the @jaxtyped decorator doesn't blat the __get__ of @classmethod
+def test_decorator():
+    assert M.f() == 3