From 81c56052e50282a3e28b3505b6640da6654cdd9e Mon Sep 17 00:00:00 2001
From: Patrick Kidger <33688385+patrick-kidger@users.noreply.github.com>
Date: Thu, 16 Feb 2023 21:08:57 +0300
Subject: [PATCH] Fixes for some new failures. (Where did they come from?)
 (#65)

* Fixes for some new failures. (Where did they come from?)

* Fixed isort?
---
 .github/workflows/run_tests.yml |  2 +-
 .pre-commit-config.yaml         |  4 ++--
 jaxtyping/__init__.py           |  2 +-
 jaxtyping/decorator.py          |  7 +------
 jaxtyping/import_hook.py        |  6 ++++--
 test/test_decorator.py          | 17 +++++++++++++++--
 6 files changed, 24 insertions(+), 14 deletions(-)

diff --git a/.github/workflows/run_tests.yml b/.github/workflows/run_tests.yml
index 56d0d23..9f1a1bf 100644
--- a/.github/workflows/run_tests.yml
+++ b/.github/workflows/run_tests.yml
@@ -26,7 +26,7 @@ jobs:
   run-tests:
     strategy:
       matrix:
-        python-version: [ 3.7, 3.8, 3.9 ]
+        python-version: [ 3.8, 3.9 ]
         os: [ ubuntu-latest ]
       fail-fast: false
     runs-on: ${{ matrix.os }}
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index e975f15..3309a43 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -23,13 +23,13 @@ repos:
     hooks:
     - id: black
   - repo: https://github.com/nbQA-dev/nbQA
-    rev: 1.2.3
+    rev: 1.6.3
     hooks:
     - id: nbqa-black
     - id: nbqa-isort 
     - id: nbqa-flake8
   - repo: https://github.com/PyCQA/isort
-    rev: 5.10.1
+    rev: 5.12.0
     hooks:
     - id: isort
   - repo: https://github.com/pycqa/flake8
diff --git a/jaxtyping/__init__.py b/jaxtyping/__init__.py
index 42ec0e7..b8feb29 100644
--- a/jaxtyping/__init__.py
+++ b/jaxtyping/__init__.py
@@ -102,4 +102,4 @@ elif has_jax:
 
 del has_jax
 
-__version__ = "0.2.11"
+__version__ = "0.2.12"
diff --git a/jaxtyping/decorator.py b/jaxtyping/decorator.py
index 41505ea..a269d00 100644
--- a/jaxtyping/decorator.py
+++ b/jaxtyping/decorator.py
@@ -31,12 +31,7 @@ class _Jaxtyped:
         self.fn = fn
 
     def __get__(self, instance, owner):
-        if instance is None:
-            # Don't create a new _Jaxtyped object in this case. Otherwise anything
-            # assigned to methods (e.g. `__isabstractmethod__`) just gets swallowed.
-            return self
-        else:
-            return ft.wraps(self.fn)(_Jaxtyped(self.fn.__get__(instance, owner)))
+        return ft.wraps(self.fn)(_Jaxtyped(self.fn.__get__(instance, owner)))
 
     def __call__(self, *args, **kwargs):
         try:
diff --git a/jaxtyping/import_hook.py b/jaxtyping/import_hook.py
index a1ed323..070817e 100644
--- a/jaxtyping/import_hook.py
+++ b/jaxtyping/import_hook.py
@@ -66,9 +66,11 @@ def _call_with_frames_removed(f, *args, **kwargs):
 
 def _optimized_cache_from_source(path, debug_override=None):
     # Version 2: change the position of the `@jaxtyped` decorator, so need a
-    # different name to avoid hitting old __pycache__
+    #     different name to avoid hitting old __pycache__.
     # Version 3: now also annotating classes.
-    return cache_from_source(path, debug_override, optimization="jaxtyping3")
+    # Version 4: I'm honestly not sure, but bumping this fixed some kind of odd error.
+    #     Maybe I changed something with hte classes part way through version 3?
+    return cache_from_source(path, debug_override, optimization="jaxtyping4")
 
 
 def _dot_lookup(*elements):
diff --git a/test/test_decorator.py b/test/test_decorator.py
index 696b25c..c9121ab 100644
--- a/test/test_decorator.py
+++ b/test/test_decorator.py
@@ -1,13 +1,26 @@
+import abc
+
 from jaxtyping import jaxtyped
 
 
-class M:
+class M(metaclass=abc.ABCMeta):
     @jaxtyped
     @classmethod
     def f(cls):
         return 3
 
+    @jaxtyped
+    @abc.abstractmethod
+    def g(self):
+        ...
+
 
 # Check that the @jaxtyped decorator doesn't blat the __get__ of @classmethod
-def test_decorator():
+def test_classmethod():
     assert M.f() == 3
+
+
+# Check that the @jaxtyped decorator doesn't blat the __isabstractmethod__ of
+# @abstractmethod
+def test_abstractmethod():
+    assert M.g.__isabstractmethod__