Fix recursive references when using class_schema()

marshmallow_dataclass/__init__.py

@@ -37,6 +37,7 @@ class User:
 import collections.abc
 import dataclasses
 import inspect
+import threading
 import types
 import warnings
 from enum import EnumMeta
@@ -77,6 +78,10 @@ class User:
 # Max number of generated schemas that class_schema keeps of generated schemas. Removes duplicates.
 MAX_CLASS_SCHEMA_CACHE_SIZE = 1024
 
+# Recursion guard for class_schema()
+_RECURSION_GUARD = threading.local()
+_RECURSION_GUARD.seen_classes = {}
+
 
 @overload
 def dataclass(
@@ -347,7 +352,10 @@ def class_schema(
             clazz_frame = current_frame.f_back
         # Per https://docs.python.org/3/library/inspect.html#the-interpreter-stack
         del current_frame
-    return _internal_class_schema(clazz, base_schema, clazz_frame)
+    try:
+        return _internal_class_schema(clazz, base_schema, clazz_frame)
+    finally:
+        _RECURSION_GUARD.seen_classes.clear()
 
 
 @lru_cache(maxsize=MAX_CLASS_SCHEMA_CACHE_SIZE)
@@ -356,6 +364,7 @@ def _internal_class_schema(
     base_schema: Optional[Type[marshmallow.Schema]] = None,
     clazz_frame: types.FrameType = None,
 ) -> Type[marshmallow.Schema]:
+    _RECURSION_GUARD.seen_classes[clazz] = clazz.__name__
     try:
         # noinspection PyDataclass
         fields: Tuple[dataclasses.Field, ...] = dataclasses.fields(clazz)
@@ -712,9 +721,11 @@ def field_for_schema(
 
     # Nested dataclasses
     forward_reference = getattr(typ, "__forward_arg__", None)
+
     nested = (
         nested_schema
         or forward_reference
+        or _RECURSION_GUARD.seen_classes.get(typ)
         or _internal_class_schema(typ, base_schema, typ_frame)
     )
 

tests/test_class_schema.py

@@ -400,6 +400,39 @@ class J:
             [validator_a, validator_b, validator_c, validator_d],
         )
 
+    def test_recursive_reference(self):
+        @dataclasses.dataclass
+        class Tree:
+            children: typing.List["Tree"]  # noqa: F821
+
+        schema = class_schema(Tree)()
+
+        self.assertEqual(
+            schema.load({"children": [{"children": []}]}),
+            Tree(children=[Tree(children=[])]),
+        )
+
+    def test_cyclic_reference(self):
+        @dataclasses.dataclass
+        class First:
+            second: typing.Optional["Second"]  # noqa: F821
+
+        @dataclasses.dataclass
+        class Second:
+            first: typing.Optional["First"]
+
+        first_schema = class_schema(First)()
+        second_schema = class_schema(Second)()
+
+        self.assertEqual(
+            first_schema.load({"second": {"first": None}}),
+            First(second=Second(first=None)),
+        )
+        self.assertEqual(
+            second_schema.dump(Second(first=First(second=Second(first=None)))),
+            {"first": {"second": {"first": None}}},
+        )
+
 
 if __name__ == "__main__":
     unittest.main()