Use WeakValueDictionary to fix generic memory leak (#6681)

dmontagu · web-flow · commit f5e3aa950ffb · 2023-07-25T12:36:57.000+01:00
diff --git a/pydantic/_internal/_model_construction.py b/pydantic/_internal/_model_construction.py
@@ -3,6 +3,7 @@
 
 import typing
 import warnings
+import weakref
 from abc import ABCMeta
 from functools import partial
 from types import FunctionType
@@ -162,8 +163,10 @@ def wrapped_model_post_init(self: BaseModel, __context: Any) -> None:
                 obj.__set_name__(cls, name)
 
             if __pydantic_reset_parent_namespace__:
-                cls.__pydantic_parent_namespace__ = parent_frame_namespace()
+                cls.__pydantic_parent_namespace__ = build_lenient_weakvaluedict(parent_frame_namespace())
             parent_namespace = getattr(cls, '__pydantic_parent_namespace__', None)
+            if isinstance(parent_namespace, dict):
+                parent_namespace = unpack_lenient_weakvaluedict(parent_namespace)
 
             types_namespace = get_cls_types_namespace(cls, parent_namespace)
             set_model_fields(cls, bases, config_wrapper, types_namespace)
@@ -560,3 +563,43 @@ def generate_model_signature(
         merged_params[var_kw_name] = var_kw.replace(name=var_kw_name)
 
     return Signature(parameters=list(merged_params.values()), return_annotation=None)
+
+
+class _PydanticWeakRef(weakref.ReferenceType):
+    pass
+
+
+def build_lenient_weakvaluedict(d: dict[str, Any] | None) -> dict[str, Any] | None:
+    """Takes an input dictionary, and produces a new value that (invertibly) replaces the values with weakrefs.
+
+    We can't just use a WeakValueDictionary because many types (including int, str, etc.) can't be stored as values
+    in a WeakValueDictionary.
+
+    The `unpack_lenient_weakvaluedict` function can be used to reverse this operation.
+    """
+    if d is None:
+        return None
+    result = {}
+    for k, v in d.items():
+        try:
+            proxy = _PydanticWeakRef(v)
+        except TypeError:
+            proxy = v
+        result[k] = proxy
+    return result
+
+
+def unpack_lenient_weakvaluedict(d: dict[str, Any] | None) -> dict[str, Any] | None:
+    """Inverts the transform performed by `build_lenient_weakvaluedict`."""
+    if d is None:
+        return None
+
+    result = {}
+    for k, v in d.items():
+        if isinstance(v, _PydanticWeakRef):
+            v = v()
+            if v is not None:
+                result[k] = v
+        else:
+            result[k] = v
+    return result
diff --git a/pydantic/main.py b/pydantic/main.py
@@ -446,10 +446,15 @@ def model_rebuild(
             else:
                 if _parent_namespace_depth > 0:
                     frame_parent_ns = _typing_extra.parent_frame_namespace(parent_depth=_parent_namespace_depth) or {}
-                    cls_parent_ns = cls.__pydantic_parent_namespace__ or {}
-                    cls.__pydantic_parent_namespace__ = {**cls_parent_ns, **frame_parent_ns}
-
-                types_namespace = cls.__pydantic_parent_namespace__
+                    cls_parent_ns = (
+                        _model_construction.unpack_lenient_weakvaluedict(cls.__pydantic_parent_namespace__) or {}
+                    )
+                    types_namespace = {**cls_parent_ns, **frame_parent_ns}
+                    cls.__pydantic_parent_namespace__ = _model_construction.build_lenient_weakvaluedict(types_namespace)
+                else:
+                    types_namespace = _model_construction.unpack_lenient_weakvaluedict(
+                        cls.__pydantic_parent_namespace__
+                    )
 
                 types_namespace = _typing_extra.get_cls_types_namespace(cls, types_namespace)
             return _model_construction.complete_model_class(
diff --git a/tests/test_generics.py b/tests/test_generics.py
@@ -405,6 +405,34 @@ class A(BaseModel, Generic[T1, T2]):
     assert len(_GENERIC_TYPES_CACHE) < types_cache_size + _LIMITED_DICT_SIZE
 
 
+@pytest.mark.skipif(platform.python_implementation() == 'PyPy', reason='PyPy does not play nice with PyO3 gc')
+def test_circular_generic_refs_get_cleaned_up():
+    initial_cache_size = len(_GENERIC_TYPES_CACHE)
+
+    def fn():
+        T = TypeVar('T')
+        C = TypeVar('C')
+
+        class Inner(BaseModel, Generic[T, C]):
+            a: T
+            b: C
+
+        class Outer(BaseModel, Generic[C]):
+            c: Inner[int, C]
+
+        klass = Outer[str]
+        assert len(_GENERIC_TYPES_CACHE) > initial_cache_size
+        assert klass in _GENERIC_TYPES_CACHE.values()
+
+    fn()
+
+    gc.collect(0)
+    gc.collect(1)
+    gc.collect(2)
+
+    assert len(_GENERIC_TYPES_CACHE) == initial_cache_size
+
+
 def test_generics_work_with_many_parametrized_base_models(clean_cache):
     cache_size = len(_GENERIC_TYPES_CACHE)
     count_create_models = 1000
