Use weak references for last context in pydantic validation (#582)

* Use weakref for previous context maintenance

Signed-off-by: Nijat Khanbabayev <nijat.khanbabayev@cubistsystematic.com>

* Add test for TVarValidationContext garbage collection

Signed-off-by: Nijat Khanbabayev <nijat.khanbabayev@cubistsystematic.com>

---------

Signed-off-by: Nijat Khanbabayev <nijat.khanbabayev@cubistsystematic.com>

Author: NeejWeej

csp/impl/types/typing_utils.py

@@ -3,6 +3,7 @@
 import sys
 import types
 import typing
+import weakref
 
 import numpy
 
@@ -178,7 +179,9 @@ def __init__(self, source_type: type):
                 f"Argument to ts must either be: a type, ForwardRef or TypeVar. Got {source_type} which is an instance of {type(source_type)}."
             )
         self._last_value_type = None
-        self._last_context = None
+        # Use a weak reference to the last validation context to avoid
+        # keeping TVarValidationContext instances alive
+        self._last_context_ref = None
 
     def validate(self, value_type, info=None):
         """Run the validation against a proposed input type"""
@@ -189,10 +192,14 @@ def validate(self, value_type, info=None):
         # is equal to the last value_type, and if so, skip validation (as any errors would already have been thrown)
         # We also don't test equality on info, assuming that the same validation info object is used
         # for a given validation run.
-        if value_type == self._last_value_type and info is not None and self._last_context is info.context:
+        last_context = None if self._last_context_ref is None else self._last_context_ref()
+        if value_type == self._last_value_type and info is not None and last_context is info.context:
             return value_type
         self._last_value_type = value_type
-        self._last_context = info.context if info is not None else None
+        if info is not None and info.context is not None:
+            self._last_context_ref = weakref.ref(info.context)
+        else:
+            self._last_context_ref = None
 
         if value_type is typing.Any:
             # https://docs.python.org/3/library/typing.html#the-any-type

csp/tests/test_engine.py

@@ -425,6 +425,26 @@ def stop(x: ts[bool]) -> ts[bool]:
         result = csp.run(stop, csp.timer(timedelta(seconds=1)), starttime=datetime(2020, 5, 19))[0]
         self.assertEqual(len(result), 5)
 
+    def test_tvar_validation_context_lifetime(self):
+        import gc
+
+        from csp.impl.types.pydantic_type_resolver import TVarValidationContext
+
+        def count_contexts():
+            return sum(1 for o in gc.get_objects() if type(o) is TVarValidationContext)
+
+        @csp.node
+        def echo(x: ts[int]) -> ts[int]:
+            return x
+
+        gc.collect(0)
+        before = count_contexts()
+        csp.build_graph(echo, realtime=False, x=csp.const(1))
+        gc.collect(0)
+        after = count_contexts()
+
+        self.assertEqual(before, after)
+
     def test_class_member_node(self):
         class ClassWithNodes:
             def __init__(self):