find static fields from dataclass parents (#14)

simple_pytree/pytree.py

@@ -239,10 +239,11 @@ def __setattr__(self: P, field: str, value: tp.Any):
 def _inherited_static_fields(cls: type) -> tp.Set[str]:
     static_fields = set()
     for parent_class in cls.mro():
-        if (
-            parent_class is not cls
-            and parent_class is not Pytree
-            and issubclass(parent_class, Pytree)
-        ):
-            static_fields.update(parent_class._pytree__static_fields)
+        if parent_class is not cls and parent_class is not Pytree:
+            if issubclass(parent_class, Pytree):
+                static_fields.update(parent_class._pytree__static_fields)
+            elif dataclasses.is_dataclass(parent_class):
+                for field in dataclasses.fields(parent_class):
+                    if not field.metadata.get("pytree_node", True):
+                        static_fields.add(field.name)
     return static_fields

tests/test_pytree.py

@@ -1,3 +1,4 @@
+import dataclasses
 from typing import Generic, TypeVar
 
 import jax
@@ -270,3 +271,21 @@ class Foo(Pytree, mutable=True):
         # test mutation
         pytree.x = 4
         assert pytree.x == 4
+
+    def test_dataclass_inheritance(self):
+        A = dataclasses.make_dataclass(
+            "A",
+            [("x", int), "y", ("z", int, static_field(default=5))],
+        )
+
+        @dataclass
+        class B(Pytree, A):
+            ...
+
+        b = B(1, 2)
+
+        assert b.x == 1
+        assert b.y == 2
+        assert b.z == 5
+
+        assert jax.tree_util.tree_leaves(b) == [1, 2]