Repository: patrick-kidger/equinox (python)

Issue:
`_wrap_method` breaks doctest's ability to find class methods
Hello!

I recently noticed that many of the doctests I wrote were never executed. After digging a bit, it seems that doctest is [not able to find methods](https://docs.python.org/3/library/doctest.html#which-docstrings-are-examined) inside subclasses of `eqx.Module`.

Unfortunately, I did not find a fix yet, but I thought it would be important to report this, and maybe you already know some fix?

E.g., the following code:

```python
import equinox as eqx


class Example(eqx.Module):
    def foo(self) -> int:
        """
        >>> example = Example()
        >>> example.foo()
        1
        """
        return 1    

    @property
    def bar(self) -> int:
        """
        >>> example = Example()
        >>> example.bar
        1
        """
        return 1

    @classmethod
    def baz(self) -> int:
        """
        >>> Example.baz()
        1
        """
        return 1

    @eqx.filter_jit
    def biz(self) -> int:
        """
        >>> example = Example()
        >>> example.biz()
        1
        """
        return 1

class Expected:
    def foo(self) -> int:
        """
        >>> example = Expected()
        >>> example.foo()
        1
        """
        return 1    

    @property
    def bar(self) -> int:
        """
        >>> example = Expected()
        >>> example.bar
        1
        """
        return 1

    @classmethod
    def baz(self) -> int:
        """
        >>> Expected.baz()
        1
        """
        return 1

if __name__ == "__main__":
    import doctest
    doctest.testmod()
```

will never test doctests inside `Example.foo`'s docstring, but all other methods will be tested.

```
(equinox-doctest) ➜  equinox-doctest git:(master) ✗ uv run python example.py -v
Cannot read termcap database;
using dumb terminal settings.
Trying:
    example = Example()
Expecting nothing
ok
Trying:
    example.bar
Expecting:
    1
ok
Trying:
    Example.baz()
Expecting:
    1
ok
Trying:
    example = Example()
Expecting nothing
ok
Trying:
    example.biz()
Expecting:
    1
ok
Trying:
    example = Expected()
Expecting nothing
ok
Trying:
    example.bar
Expecting:
    1
ok
Trying:
    Expected.baz()
Expecting:
    1
ok
Trying:
    example = Expected()
Expecting nothing
ok
Trying:
    example.foo()
Expecting:
    1
ok
6 items had no tests:
    __main__
    __main__.Example
    __main__.Example.__delattr__
    __main__.Example.__init__
    __main__.Example.__setattr__
    __main__.Expected
6 items passed all tests:
   2 tests in __main__.Example.bar
   1 tests in __main__.Example.baz
   2 tests in __main__.Example.biz
   2 tests in __main__.Expected.bar
   1 tests in __main__.Expected.baz
   2 tests in __main__.Expected.foo
10 tests in 12 items.
10 passed and 0 failed.
Test passed.
```


Pull request:
fix(lib): allow DocTest to discover methods from subclasses `Module`
Closes #1016.

In the end, there were two issues: (1) `DocTestFinder._from_module` was returning `False` because the `__module__` attribute was not correct, and (2) the documentation was not copied to the wrapper class so doctests could not be discovered.

My solution is inspired from what Python is doing with their own decorators, see https://github.com/python/cpython/blob/4617d68d73409e83d6ab31106d10421d44048787/Lib/functools.py#L1033-L1049

Patch:
diff --git a/equinox/_module.py b/equinox/_module.py
index 69a1bc7..a67c8e3 100644
--- a/equinox/_module.py
+++ b/equinox/_module.py
@@ -681,6 +681,15 @@ class _wrap_method:
         if getattr(self.method, "__isabstractmethod__", False):
             self.__isabstractmethod__ = self.method.__isabstractmethod__
 
+    # Fixes https://github.com/patrick-kidger/equinox/issues/1016
+    @property
+    def __module__(self):
+        return self.method.__module__
+
+    @property
+    def __doc__(self):
+        return self.method.__doc__
+
     def __get__(self, instance, owner):
         del owner
         if instance is None:


Test patch:
diff --git a/.github/workflows/run_tests.yml b/.github/workflows/run_tests.yml
index 0934f36..806ae01 100644
--- a/.github/workflows/run_tests.yml
+++ b/.github/workflows/run_tests.yml
@@ -27,7 +27,7 @@ jobs:
           python -m pip install -r ./tests/requirements.txt
 
       - name: Checks with pre-commit
-        uses: pre-commit/action@v2.0.3
+        uses: pre-commit/action@v3.0.1
 
       - name: Test with pytest
         run: |
diff --git a/tests/test_module.py b/tests/test_module.py
index b9a7471..680341d 100644
--- a/tests/test_module.py
+++ b/tests/test_module.py
@@ -1,5 +1,6 @@
 import abc
 import dataclasses
+import doctest
 import functools as ft
 import inspect
 import pickle
@@ -1257,3 +1258,61 @@ def test_tree_roundtrip_serialise_roundtrip():
     x = pickle.loads(pickle.dumps(x))
     tree2 = jtu.tree_structure(x)
     assert tree == tree2
+
+
+# https://github.com/patrick-kidger/equinox/issues/1016
+def test_doctest():
+    class Example(eqx.Module):
+        def foo(self) -> int:
+            """
+            >>> example = Example()
+            >>> example.foo()
+            1
+            """
+            return 1
+
+        @property
+        def bar(self) -> int:
+            """
+            >>> example = Example()
+            >>> example.bar
+            1
+            """
+            return 1
+
+        @classmethod
+        def baz(self) -> int:
+            """
+            >>> Example.baz()
+            1
+            """
+            return 1
+
+        @eqx.filter_jit
+        def biz(self) -> int:
+            """
+            >>> example = Example()
+            >>> example.biz()
+            1
+            """
+            return 1
+
+        @staticmethod
+        def buz() -> int:
+            """
+            >>> example = Example()
+            >>> example.buz()
+            1
+            """
+            return 1
+
+    tests = doctest.DocTestFinder().find(Example)
+    tests = [test for test in tests if test.examples]
+
+    assert len(tests) == 5
+
+    assert any("foo" in str(test) for test in tests)
+    assert any("bar" in str(test) for test in tests)
+    assert any("baz" in str(test) for test in tests)
+    assert any("biz" in str(test) for test in tests)
+    assert any("buz" in str(test) for test in tests)