Repository: beartype/plum

Issue:
`parametric` conflicts with certain usages of customized `__init_subclass__`
Hi!

I'm using JAX, and also using `plum` -- in my library, I've define a mixin class called `Pytree` which automatically implements the Pytree interface for classes which mix it in.

It's quite simple:

```python
class Pytree:
    def __init_subclass__(cls, **kwargs):
        jtu.register_pytree_node(
            cls,
            cls.flatten,
            cls.unflatten,
        )

```

If I wish to use this mixin, and `parametric` -- I'm in for problems, I get duplicate registration:

```
ERROR ... ValueError: Duplicate custom PyTreeDef type registration for <class...>
```

I'm not exactly sure why this occurs, but I'm hoping to find a fix -- because I'd like to use parametric classes to guide some of the dispatch in my library functions.

Comments:

Patch:
diff --git a/plum/function.py b/plum/function.py
index 7b97c00..15b4b9c 100644
--- a/plum/function.py
+++ b/plum/function.py
@@ -1,3 +1,4 @@
+import os
 import textwrap
 from functools import wraps
 from types import MethodType
@@ -141,6 +142,11 @@ class Function(metaclass=_FunctionMeta):
             # clearing the cache.
             self.clear_cache(reregister=False)
 
+        # Don't do any fancy appending of docstrings when the environment variable
+        # `PLUM_SIMPLE_DOC` is set to `1`.
+        if "PLUM_SIMPLE_DOC" in os.environ and os.environ["PLUM_SIMPLE_DOC"] == "1":
+            return self._doc
+
         # Derive the basis of the docstring from `self._f`, removing any indentation.
         doc = self._doc.strip()
         if doc:
diff --git a/plum/parametric.py b/plum/parametric.py
index ec69cf3..12aafa3 100644
--- a/plum/parametric.py
+++ b/plum/parametric.py
@@ -240,7 +240,7 @@ def parametric(original_class=None):
                 return original_class.__new__(cls)
 
             cls.__new__ = class_new
-        original_class.__init_subclass__(**kw_args)
+        super(original_class, cls).__init_subclass__(**kw_args)
 
     # Create parametric class.
     parametric_class = meta(
diff --git a/pyproject.toml b/pyproject.toml
index 6671c47..addea61 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -19,7 +19,7 @@ dynamic = ["version"]
 
 requires-python = ">=3.8"
 dependencies = [
-    "beartype",
+    "beartype >= 0.16",
     "typing-extensions; python_version<='3.10'",
 ]
 


Test patch:
diff --git a/tests/test_function.py b/tests/test_function.py
index c4f8229..234875d 100644
--- a/tests/test_function.py
+++ b/tests/test_function.py
@@ -1,4 +1,5 @@
 import abc
+import os
 import textwrap
 import typing
 
@@ -221,6 +222,31 @@ def test_doc(monkeypatch):
     assert g.__doc__ == textwrap.dedent(expected_doc).strip()
 
 
+def test_simple_doc(monkeypatch):
+    @dispatch
+    def f(x: int):
+        """First."""
+
+    @dispatch
+    def f(x: str):
+        """Second."""
+
+    monkeypatch.setitem(os.environ, "PLUM_SIMPLE_DOC", "1")
+    assert f.__doc__ == "First."
+
+    monkeypatch.setitem(os.environ, "PLUM_SIMPLE_DOC", "0")
+    expected_doc = """
+    First.
+
+    -----------
+
+    f(x: str)
+
+    Second.
+    """
+    assert f.__doc__ == textwrap.dedent(expected_doc).strip()
+
+
 def test_methods():
     dispatch = Dispatcher()
 
diff --git a/tests/test_parametric.py b/tests/test_parametric.py
index d9fa79c..00bbcfa 100644
--- a/tests/test_parametric.py
+++ b/tests/test_parametric.py
@@ -507,3 +507,25 @@ def test_val():
         Val[1].__init__(MockVal())
 
     assert repr(Val[1]()) == "plum.parametric.Val[1]()"
+
+
+def test_init_subclass_correct_args():
+    # See issue https://github.com/beartype/plum/issues/105
+
+    from plum import parametric
+
+    register = set()
+
+    class Pytree:
+        def __init_subclass__(cls, **kwargs):
+            if cls in register:
+                raise ValueError("duplicate")
+            else:
+                register.add(cls)
+
+    @parametric
+    class Wrapper(Pytree):
+        pass
+
+    Wrapper[int]
+    assert Wrapper[int] in register