diff --git a/jaxtyping/decorator.py b/jaxtyping/decorator.py index 3100904..9c329b4 100644 --- a/jaxtyping/decorator.py +++ b/jaxtyping/decorator.py @@ -21,35 +21,20 @@ import dataclasses import functools as ft import inspect import threading +import types +import weakref storage = threading.local() -class _Jaxtyped: - def __get__(self, instance, owner): - fn = self.__wrapped__ - got = fn.__get__(instance, owner) - if fn is got: - return self - else: - return ft.wraps(got)(_Jaxtyped()) - - def __call__(self, *args, **kwargs): - try: - memo_stack = storage.memo_stack - except AttributeError: - memo_stack = storage.memo_stack = [] - memo_stack.append(({}, {}, {})) - fn = self.__wrapped__ - try: - return fn(*args, **kwargs) - finally: - memo_stack.pop() +_jaxtyped_fns = weakref.WeakSet() def jaxtyped(fn): - if inspect.isclass(fn): # allow decorators on class definitions + if type(fn) is types.FunctionType and fn in _jaxtyped_fns: + return fn + elif inspect.isclass(fn): # allow decorators on class definitions if dataclasses.is_dataclass(fn): init = jaxtyped(fn.__init__) fn.__init__ = init @@ -58,8 +43,44 @@ def jaxtyped(fn): raise ValueError( "jaxtyped may only be added as a class decorator to dataclasses" ) + # It'd be lovely if we could handle arbitrary descriptors, and not just the builtin + # ones. Unfortunately that means returning a class instance with a __get__ method, + # and that turns out to break loads of other things. See beartype issue #211 and + # jaxtyping issue #71. + elif isinstance(fn, classmethod): + return classmethod(jaxtyped(fn.__func__)) + elif isinstance(fn, staticmethod): + return staticmethod(jaxtyped(fn.__func__)) + elif isinstance(fn, property): + if fn.fget is None: + fget = None + else: + fget = jaxtyped(fn.fget) + if fn.fset is None: + fset = None + else: + fset = jaxtyped(fn.fset) + if fn.fdel is None: + fdel = None + else: + fdel = jaxtyped(fn.fdel) + return property(fget=fget, fset=fset, fdel=fdel) else: - return ft.wraps(fn)(_Jaxtyped()) + + @ft.wraps(fn) + def wrapped_fn(*args, **kwargs): + try: + memo_stack = storage.memo_stack + except AttributeError: + memo_stack = storage.memo_stack = [] + memo_stack.append(({}, {}, {})) + try: + return fn(*args, **kwargs) + finally: + memo_stack.pop() + + _jaxtyped_fns.add(wrapped_fn) + return wrapped_fn def _jaxtyped_typechecker(typechecker): @@ -72,10 +93,8 @@ def _jaxtyped_typechecker(typechecker): def _wrapper(kls): assert inspect.isclass(kls) if dataclasses.is_dataclass(kls): - if type(kls.__init__) is not _Jaxtyped: - # Extra `if` check to work around beartype bug #211 - init = jaxtyped(typechecker(kls.__init__)) - kls.__init__ = init + init = jaxtyped(typechecker(kls.__init__)) + kls.__init__ = init return kls return _wrapper diff --git a/test/import_hook_tester_typeguard.py b/test/import_hook_tester_typeguard.py index b77b756..b184215 100644 --- a/test/import_hook_tester_typeguard.py +++ b/test/import_hook_tester_typeguard.py @@ -17,9 +17,7 @@ # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. -import dataclasses -import equinox as eqx import jax.numpy as jnp import pytest @@ -37,26 +35,29 @@ with pytest.raises(ParamError): g(jnp.array(1)) -class M(eqx.Module): - foo: int - bar: Float32[jnp.ndarray, " a"] - - -M(1, jnp.array([1.0])) -with pytest.raises(ParamError): - M(1.0, jnp.array([1.0])) -with pytest.raises(ParamError): - M(1, jnp.array(1.0)) - - -@dataclasses.dataclass -class D: - foo: int - bar: Float32[jnp.ndarray, " a"] - - -D(1, jnp.array([1.0])) -with pytest.raises(ParamError): - D(1.0, jnp.array([1.0])) -with pytest.raises(ParamError): - D(1, jnp.array(1.0)) +# Typeguard 3.0 no longer supports this. +# +# class M(eqx.Module): +# foo: int +# bar: Float32[jnp.ndarray, " a"] +# +# +# M(1, jnp.array([1.0])) +# with pytest.raises(ParamError): +# M(1.0, jnp.array([1.0])) +# with pytest.raises(ParamError): +# M(1, jnp.array(1.0)) +# +# +# +# @dataclasses.dataclass +# class D: +# foo: int +# bar: Float32[jnp.ndarray, " a"] +# +# +# D(1, jnp.array([1.0])) +# with pytest.raises(ParamError): +# D(1.0, jnp.array([1.0])) +# with pytest.raises(ParamError): +# D(1, jnp.array(1.0)) diff --git a/test/import_hook_tester_typeguard_old.py b/test/import_hook_tester_typeguard_old.py index b77b756..b184215 100644 --- a/test/import_hook_tester_typeguard_old.py +++ b/test/import_hook_tester_typeguard_old.py @@ -17,9 +17,7 @@ # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. -import dataclasses -import equinox as eqx import jax.numpy as jnp import pytest @@ -37,26 +35,29 @@ with pytest.raises(ParamError): g(jnp.array(1)) -class M(eqx.Module): - foo: int - bar: Float32[jnp.ndarray, " a"] - - -M(1, jnp.array([1.0])) -with pytest.raises(ParamError): - M(1.0, jnp.array([1.0])) -with pytest.raises(ParamError): - M(1, jnp.array(1.0)) - - -@dataclasses.dataclass -class D: - foo: int - bar: Float32[jnp.ndarray, " a"] - - -D(1, jnp.array([1.0])) -with pytest.raises(ParamError): - D(1.0, jnp.array([1.0])) -with pytest.raises(ParamError): - D(1, jnp.array(1.0)) +# Typeguard 3.0 no longer supports this. +# +# class M(eqx.Module): +# foo: int +# bar: Float32[jnp.ndarray, " a"] +# +# +# M(1, jnp.array([1.0])) +# with pytest.raises(ParamError): +# M(1.0, jnp.array([1.0])) +# with pytest.raises(ParamError): +# M(1, jnp.array(1.0)) +# +# +# +# @dataclasses.dataclass +# class D: +# foo: int +# bar: Float32[jnp.ndarray, " a"] +# +# +# D(1, jnp.array([1.0])) +# with pytest.raises(ParamError): +# D(1.0, jnp.array([1.0])) +# with pytest.raises(ParamError): +# D(1, jnp.array(1.0)) diff --git a/test/test_decorator.py b/test/test_decorator.py index 25d0bdb..c655425 100644 --- a/test/test_decorator.py +++ b/test/test_decorator.py @@ -4,43 +4,74 @@ from jaxtyping import jaxtyped class M(metaclass=abc.ABCMeta): + @jaxtyped + def f(self): + ... + @jaxtyped @classmethod - def f1(cls): + def g1(cls): return 3 @classmethod @jaxtyped - def f2(cls): + def g2(cls): + return 4 + + @jaxtyped + @staticmethod + def h1(): + return 3 + + @staticmethod + @jaxtyped + def h2(): return 4 @jaxtyped @abc.abstractmethod - def g1(self): + def i1(self): ... @abc.abstractmethod @jaxtyped - def g2(self): + def i2(self): ... + +class N: @jaxtyped - def h(self): - ... + @property + def j1(self): + return 3 + + @property + @jaxtyped + def j2(self): + return 4 + + +def test_identity(): + assert M.f is M.f -# Check that the @jaxtyped decorator doesn't blat the __get__ of @classmethod def test_classmethod(): - assert M.f1() == 3 - assert M.f2() == 4 + assert M.g1() == 3 + assert M.g2() == 4 + + +def test_staticmethod(): + assert M.h1() == 3 + assert M.h2() == 4 # Check that the @jaxtyped decorator doesn't blat the __isabstractmethod__ of # @abstractmethod def test_abstractmethod(): - assert M.g1.__isabstractmethod__ - assert M.g2.__isabstractmethod__ + assert M.i1.__isabstractmethod__ + assert M.i2.__isabstractmethod__ -def test_identity(): - assert M.h is M.h +def test_property(): + assert N().j1 == 3 + assert N().j2 == 4 diff --git a/test/test_import_hook.py b/test/test_import_hook.py index 9f35fe4..cff30df 100644 --- a/test/test_import_hook.py +++ b/test/test_import_hook.py @@ -78,7 +78,7 @@ def test_import_hook_beartype_full(): def test_import_hook_transitive(): hook = install_import_hook( - "test.import_hook_tester_transitive", "typeguard.typechecked" + "test.import_hook_tester_transitive", "beartype.beartype" ) with hook: from . import import_hook_tester_transitive # noqa: F401