From 79f9f7459fd4453720d3c4f7232abf28c09a4814 Mon Sep 17 00:00:00 2001
From: Paul Prescod <pprescod@salesforce.com>
Date: Mon, 1 Aug 2022 01:54:49 -0230
Subject: [PATCH] Close plugins during Interpreter clean-up

---
 snowfakery/data_generator_runtime.py       |  4 ++
 snowfakery/plugins.py                      | 24 ++++++--
 tests/test_custom_plugins_and_providers.py | 70 +++++++++++++++++++++-
 3 files changed, 93 insertions(+), 5 deletions(-)

diff --git a/snowfakery/data_generator_runtime.py b/snowfakery/data_generator_runtime.py
index cad83c88..bcbaf02c 100644
--- a/snowfakery/data_generator_runtime.py
+++ b/snowfakery/data_generator_runtime.py
@@ -440,6 +440,10 @@ def __exit__(self, *args):
                 plugin.close()
             except Exception as e:
                 warn(f"Could not close {plugin} because {e}")
+        self.current_context = None
+        self.plugin_instances = None
+        self.plugin_function_libraries = None
+        self.instance_states = None
 
     def get_contextual_state(
         self,
diff --git a/snowfakery/plugins.py b/snowfakery/plugins.py
index c7c548ce..88b7780b 100644
--- a/snowfakery/plugins.py
+++ b/snowfakery/plugins.py
@@ -289,16 +289,25 @@ def __init__(self, result: Mapping):
 
     def __getattr__(self, name):
         # ensures that it won't recurse
-        return self.__dict__["result"][name]
+        res = self.__dict__.get("result", {}).get(name, ...)
+        if res == ...:
+            raise AttributeError(name)
 
     def __reduce__(self):
         return (self.__class__, (dict(self.result),))
 
     def __repr__(self):
-        return f"<{self.__class__} {repr(self.result)}>"
+        try:
+            rep = repr(self.result)
+        except Exception:
+            rep = ""
+        return f"<{self.__class__} {rep}>"
 
     def __str__(self):
-        return str(self.result)
+        try:
+            return str(self.result)
+        except Exception:
+            return repr(self)
 
     @classmethod
     def _from_continuation(cls, args):
@@ -320,6 +329,8 @@ def _register_for_continuation(cls):
 
 
 class PluginResultIterator(PluginResult):
+    closed = False
+
     def __init__(self, repeat):
         self.repeat = repeat
 
@@ -353,7 +364,12 @@ def restart(self):
 
     def close(self):
         "Subclasses should implement this if they need to clean up resources"
-        pass  # pragma: no cover
+        pass
+
+    def __del__(self):
+        if not self.closed:
+            self.close()
+            self.close = True
 
 
 class PluginOption:
diff --git a/tests/test_custom_plugins_and_providers.py b/tests/test_custom_plugins_and_providers.py
index d26b5fb2..179f164c 100644
--- a/tests/test_custom_plugins_and_providers.py
+++ b/tests/test_custom_plugins_and_providers.py
@@ -3,9 +3,15 @@
 import operator
 import time
 from base64 import b64decode
+import gc
 
 from snowfakery import SnowfakeryPlugin, lazy
-from snowfakery.plugins import PluginResult, PluginOption, memorable
+from snowfakery.plugins import (
+    PluginResult,
+    PluginOption,
+    PluginResultIterator,
+    memorable,
+)
 
 from snowfakery.data_gen_exceptions import (
     DataGenError,
@@ -40,6 +46,45 @@ class Functions:
         def double(self, value):
             return value * 2
 
+        @memorable
+        def fib(self, value):
+            return FibIterator()
+
+
+class FibIterator(PluginResultIterator):
+    def __init__(self):
+        self.a = 1
+        self.b = 1
+
+    def close(self):
+        self.a = None
+        self.b = None
+
+    def next(self):
+        # import gc
+
+        # for ref in gc.get_referrers(self):
+        #     print(ref, gc.get_referrers(ref))
+        # print(vars(self))
+        # print(self in gc.get_objects())
+        # print_referrers_bredth_first((self,), 0)
+        self.a, self.b = self.b, self.a + self.b
+        return self.b
+
+
+def print_referrers_bredth_first(path, level):
+    if level > 3:
+        return
+    obj = path[0]
+    parents = tuple(gc.get_referrers(obj))
+    for parent in parents:
+        parts = (parent,) + path
+        print(" -> ".join(repr(part) for part in parts))
+
+    for parent in parents:
+        parts = (parent,) + path
+        print_referrers_bredth_first(parts, level + 1)
+
 
 class DoubleVisionPlugin(SnowfakeryPlugin):
     class Functions:
@@ -414,3 +459,26 @@ def test_plugin_does_not_close(self):
         """
         with pytest.warns(UserWarning, match="close"):
             generate_data(StringIO(yaml))
+
+
+class TestPluginResultIterator:
+    @mock.patch(
+        "tests.test_custom_plugins_and_providers.FibIterator.close",
+    )
+    def test_plugin_result_iterator__closes(self, close, generated_rows):
+        yaml = """
+        - plugin: tests.test_custom_plugins_and_providers.SimpleTestPlugin
+        - object: OBJ
+          fields:
+            fibnum:
+                SimpleTestPlugin.fib:
+        """
+        # it's better if closing of objects is triggered at a predictable
+        # time by the refcounter instead of by the cyclic GC
+        try:
+            gc.disable()
+            generate_data(StringIO(yaml), target_number=("OBJ", 3))
+            assert generated_rows.table_values("OBJ", field="fibnum") == [2, 3, 5]
+            assert len(close.mock_calls) == 1
+        finally:
+            gc.enable()