From 784aa78f7c282c83e578f39b3fe4ea4a149df726 Mon Sep 17 00:00:00 2001
From: Kevin P Murphy <murphyk@gmail.com>
Date: Mon, 14 Nov 2022 07:25:28 +0100
Subject: [PATCH] update jaxtyped decorator (#44)

* update jaxtyped decorator

* add newline character to pacify flake

* add precomit hooks

* add missing return statement

* replace typing_extensions>=4.2.0 with typing_extensions

* pin version range for typing_extensions

* set min version of typing-extensions but not max

* bump version number to 0.2.8
---
 jaxtyping/__init__.py  | 2 +-
 jaxtyping/decorator.py | 8 +++++++-
 setup.py               | 5 ++++-
 3 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/jaxtyping/__init__.py b/jaxtyping/__init__.py
index cf9441c..932a5f9 100644
--- a/jaxtyping/__init__.py
+++ b/jaxtyping/__init__.py
@@ -67,4 +67,4 @@ from .import_hook import install_import_hook
 from .pytree_type import PyTree
 
 
-__version__ = "0.2.7"
+__version__ = "0.2.8"
diff --git a/jaxtyping/decorator.py b/jaxtyping/decorator.py
index b4916f5..61d82bb 100644
--- a/jaxtyping/decorator.py
+++ b/jaxtyping/decorator.py
@@ -18,6 +18,7 @@
 # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
 import functools as ft
+import inspect
 import threading
 
 
@@ -44,4 +45,9 @@ class _Jaxtyped:
 
 
 def jaxtyped(fn):
-    return ft.wraps(fn)(_Jaxtyped(fn))
+    if inspect.isclass(fn):  # allow decorators on class definitions
+        init = jaxtyped(fn.__init__)
+        fn.__init__ = init
+        return fn
+    else:
+        return ft.wraps(fn)(_Jaxtyped(fn))
diff --git a/setup.py b/setup.py
index 7185d91..6541801 100644
--- a/setup.py
+++ b/setup.py
@@ -67,11 +67,14 @@ python_requires = "~=3.7"
 
 # We use typeguard internally (in a fairly minimal way), but it's not required that
 # end users make the same choice.
+# For typing_extensions, we choose versions that match
+# https://github.com/explosion/confection/blob/main/setup.cfg#L33 used in colab
+
 install_requires = [
     "jax>=0.3.4",
     "numpy>=1.20.0",
     "typeguard>=2.13.3",
-    "typing_extensions>=4.2.0",
+    "typing_extensions>=3.7.4.1",
 ]
 
 entry_points = dict(pytest11=["jaxtyping = jaxtyping.pytest_plugin"])