Apply functools.wraps to Module methods to preserve docstrings

patrick-kidger · Jul 5, 2023 · 948b633 · 948b633
1 parent b7c6f63
commit 948b633
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/equinox/_module.py b/equinox/_module.py
@@ -101,7 +101,9 @@ def __init__(self, method):
     def __get__(self, instance, owner):
         if instance is None:
             return self.method
-        return jtu.Partial(self.method, instance)
+        _method = ft.wraps(self.method)(jtu.Partial(self.method, instance))
+        delattr(_method, "__wrapped__")
+        return _method
 
 
 def _not_magic(k: str) -> bool: