PyTrees now use Wadler-Lindig to handle their naming and pretty-printing. In particualr this is useful for better documentation.

Author: patrick-kidger

jaxtyping/_pytree_type.py

@@ -22,6 +22,7 @@
 from typing import Any, Generic, TypeVar
 
 import jax.tree_util as jtu
+import wadler_lindig as wl
 
 from ._errors import AnnotationError
 from ._storage import (
@@ -35,15 +36,25 @@
 
 
 _T = TypeVar("_T")
+_S = TypeVar("_S")
 
 
-class _FakePyTree(Generic[_T]):
+class _FakePyTree1(Generic[_T]):
     pass
 
 
-_FakePyTree.__name__ = "PyTree"
-_FakePyTree.__qualname__ = "PyTree"
-_FakePyTree.__module__ = "builtins"
+_FakePyTree1.__name__ = "PyTree"
+_FakePyTree1.__qualname__ = "PyTree"
+_FakePyTree1.__module__ = "builtins"
+
+
+class _FakePyTree2(Generic[_T, _S]):
+    pass
+
+
+_FakePyTree2.__name__ = "PyTree"
+_FakePyTree2.__qualname__ = "PyTree"
+_FakePyTree2.__module__ = "builtins"
 
 
 class _MetaPyTree(type):
@@ -226,7 +237,15 @@ class X(PyTree):
                             "regular Python, i.e. a valid variable name.)\n"
                             f"Got piece '{piece}' in overall structure '{X.structure}'."
                         )
-                name = str(_FakePyTree[item[0]])[:-1] + ', "' + item[1].strip() + '"]'
+
+                class Y:
+                    pass
+
+                Y.__module__ = "builtins"
+                Y.__name__ = repr(X.structure)
+                Y.__qualname__ = repr(X.structure)
+                name = wl.pformat(_FakePyTree2[X.leaftype, Y], width=9999)
+                del Y
             else:
                 raise ValueError(
                     "The subscript `foo` in `jaxtyping.PyTree[foo]` must either be a "
@@ -235,7 +254,7 @@ class X(PyTree):
                     f"{len(item)}."
                 )
         else:
-            name = str(_FakePyTree[item])
+            name = wl.pformat(_FakePyTree1[item], width=9999)
 
             class X(PyTree):
                 leaftype = item
@@ -249,6 +268,22 @@ class X(PyTree):
             X.__module__ = "jaxtyping"
         return X
 
+    def __pdoc__(self, **kwargs):
+        if self is PyTree:
+            return wl.TextDoc("PyTree")
+        else:
+            indent = kwargs["indent"]
+            docs = [wl.pdoc(self.leaftype, **kwargs)]
+            if self.structure is not None:
+                docs.append(wl.pdoc(self.structure, **kwargs))
+            return wl.bracketed(
+                begin=wl.TextDoc("PyTree["),
+                docs=docs,
+                sep=wl.comma,
+                end=wl.TextDoc("]"),
+                indent=indent,
+            )
+
 
 # Can't do `class PyTree(Generic[_T]): ...` because we need to override the
 # instancecheck for PyTree[foo], but subclassing

test/test_pytree.py

@@ -17,13 +17,15 @@
 # IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 # CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
+from collections.abc import Callable
 from typing import NamedTuple, Tuple, Union
 
 import equinox as eqx
 import jax
 import jax.numpy as jnp
 import jax.random as jr
 import pytest
+import wadler_lindig as wl
 
 import jaxtyping
 from jaxtyping import AnnotationError, Array, Float, PyTree
@@ -341,3 +343,53 @@ def f(x: PyTree[PyTree[Float[Array, "?foo"], " S"], " T"]):
     x1 = jr.normal(getkey(), (2,))
     with pytest.raises(AnnotationError, match="ambiguous which PyTree"):
         f(x1)
+
+
+def test_name():
+    assert PyTree.__name__ == "PyTree"
+    assert PyTree[int].__name__ == "PyTree[int]"
+    assert PyTree[int, "foo"].__name__ == "PyTree[int, 'foo']"
+    assert PyTree[PyTree[str], "foo"].__name__ == "PyTree[PyTree[str], 'foo']"
+    assert (
+        PyTree[PyTree[str, "bar"], "foo"].__name__
+        == "PyTree[PyTree[str, 'bar'], 'foo']"
+    )
+    assert PyTree[PyTree[str, "bar"]].__name__ == "PyTree[PyTree[str, 'bar']]"
+    assert (
+        PyTree[None | Callable[[PyTree[int, " T"]], str]].__name__
+        == "PyTree[None | Callable[[PyTree[int, 'T']], str]]"
+    )
+
+
+def test_pdoc():
+    assert wl.pformat(PyTree) == "PyTree"
+    assert wl.pformat(PyTree[int]) == "PyTree[int]"
+    assert wl.pformat(PyTree[int, "foo"]) == "PyTree[int, 'foo']"
+    assert wl.pformat(PyTree[PyTree[str], "foo"]) == "PyTree[PyTree[str], 'foo']"
+    assert (
+        wl.pformat(PyTree[PyTree[str, "bar"], "foo"])
+        == "PyTree[PyTree[str, 'bar'], 'foo']"
+    )
+    assert wl.pformat(PyTree[PyTree[str, "bar"]]) == "PyTree[PyTree[str, 'bar']]"
+    assert (
+        wl.pformat(PyTree[None | Callable[[PyTree[int, " T"]], str]])
+        == "PyTree[None | Callable[[PyTree[int, 'T']], str]]"
+    )
+    expected = """
+PyTree[
+  None
+  | Callable[
+    [
+      PyTree[
+        int,
+        'T'
+      ]
+    ],
+    str
+  ]
+]
+    """.strip()
+    assert (
+        wl.pformat(PyTree[None | Callable[[PyTree[int, " T"]], str]], width=2).strip()
+        == expected
+    )