Support recursive and empty closure cells

dill/_dill.py

@@ -481,6 +481,7 @@ def __init__(self, *args, **kwds):
         self._strictio = False #_strictio
         self._fmode = settings['fmode'] if _fmode is None else _fmode
         self._recurse = settings['recurse'] if _recurse is None else _recurse
+        self._recursive_cells = {}
 
     def dump(self, obj): #NOTE: if settings change, need to update attributes
         stack.clear()  # clear record of 'recursion-sensitive' pickled objects
@@ -879,9 +880,31 @@ def __getattribute__(self, attr):
 if PY3:
     def _create_cell(contents):
         return (lambda y: contents).__closure__[0]
+    exec('''
+def _create_reference_cell():
+    contents = None
+    def updater(value):
+        nonlocal contents
+        contents = value
+    updater(updater)
+    return (lambda: contents).__closure__[0]
+''')
 else:
     def _create_cell(contents):
         return (lambda y: contents).func_closure[0]
+    def _create_reference_cell():
+        contents = None
+        v = vars()
+        class Updater(object):
+            def __call__(self, value):
+                v['contents'] = value
+        contents = Updater()
+        return (lambda: contents).func_closure[0]
+
+    # _create_reference_cell not possible in Python 2
+
+def _update_cell(cell, obj_ptr):
+    return cell.cell_contents(obj_ptr)
 
 def _create_weakref(obj, *args):
     from weakref import ref
@@ -1296,10 +1319,18 @@ def save_wrapper_descriptor(pickler, obj):
 
 @register(CellType)
 def save_cell(pickler, obj):
-    log.info("Ce: %s" % obj)
     f = obj.cell_contents
+    if is_dill(pickler, child=True):
+        recursive_cells = pickler._recursive_cells.get(id(f))
+        if recursive_cells is not None:
+            log.info("Ce2: %s" % obj)
+            pickler.save_reduce(_create_reference_cell, (), obj=obj)
+            recursive_cells.append(obj)
+            log.info("# Ce2")
+            return
+    log.info("Ce1: %s" % obj)
     pickler.save_reduce(_create_cell, (f,), obj=obj)
-    log.info("# Ce")
+    log.info("# Ce1")
     return
 
 if not IS_PYPY:
@@ -1470,9 +1501,10 @@ def save_type(pickler, obj):
         log.info("# T6")
         return
     elif obj.__module__ == '__main__':
+        pickler_is_dill = is_dill(pickler, child=True)
         if issubclass(type(obj), type):
         #   try: # used when pickling the class as code (or the interpreter)
-            if is_dill(pickler, child=True) and not pickler._byref:
+            if pickler_is_dill and not pickler._byref and id(obj) not in pickler._recursive_cells:
                 # thanks to Tom Stepleton pointing out pickler._session unneeded
                 _t = 'T2'
                 log.info("%s: %s" % (_t, obj))
@@ -1493,8 +1525,20 @@ def save_type(pickler, obj):
        #print ("%s\n%s" % (obj.__bases__, obj.__dict__))
         for name in _dict.get("__slots__", []):
             del _dict[name]
-        pickler.save_reduce(_create_type, (type(obj), obj.__name__,
+        if pickler_is_dill:
+            pickler._recursive_cells[id(obj)] = []
+        name = getattr(obj, "__qualname__", obj.__name__)
+        pickler.save_reduce(_create_type, (type(obj), name,
                                            obj.__bases__, _dict), obj=obj)
+        if pickler_is_dill:
+            recursive_cells = pickler._recursive_cells.pop(id(obj))
+            for t in recursive_cells:
+                pickler.save_reduce(_update_cell, (t, obj))
+                # pop None off created by setattr off stack
+                if PY3:
+                    pickler.write(bytes('0', 'UTF-8'))
+                else:
+                    pickler.write('0')
         log.info("# %s" % _t)
     # special cases: NoneType, NotImplementedType, EllipsisType
     elif obj is type(None):
@@ -1570,7 +1614,6 @@ def save_function(pickler, obj):
         if PY3:
             #NOTE: workaround for 'super' (see issue #75)
             _super = ('super' in getattr(obj.__code__,'co_names',())) and (_byref is not None)
-            if _super: pickler._byref = True
             if _memo: pickler._recurse = False
             fkwdefaults = getattr(obj, '__kwdefaults__', None)
             pickler.save_reduce(_create_function, (obj.__code__,
@@ -1585,7 +1628,7 @@ def save_function(pickler, obj):
                                 globs, obj.func_name,
                                 obj.func_defaults, obj.func_closure,
                                 obj.__dict__), obj=obj)
-        if _super: pickler._byref = _byref
+            if _super: pickler._byref = _byref
         if _memo: pickler._recurse = _recurse
        #clear = (_byref, _super, _recurse, _memo)
        #print(clear + (OLDER,))

tests/test_recursive.py

@@ -27,17 +27,17 @@ def __init__(self):
 def test_super():
     assert dill.copy(obj1(), byref=True)
     assert dill.copy(obj1(), byref=True, recurse=True)
-   #assert dill.copy(obj1(), recurse=True) #FIXME: fails __main__.py
+    assert dill.copy(obj1(), recurse=True)
     assert dill.copy(obj1())
 
     assert dill.copy(obj2(), byref=True)
     assert dill.copy(obj2(), byref=True, recurse=True)
-   #assert dill.copy(obj2(), recurse=True) #FIXME: fails __main__.py
+    assert dill.copy(obj2(), recurse=True)
     assert dill.copy(obj2())
 
     assert dill.copy(obj3(), byref=True)
     assert dill.copy(obj3(), byref=True, recurse=True)
-   #assert dill.copy(obj3(), recurse=True) #FIXME: fails __main__.py
+    assert dill.copy(obj3(), recurse=True)
     assert dill.copy(obj3())
 
 
@@ -58,8 +58,7 @@ class Model(object):
 def test_partial():
     assert dill.copy(Machine(), byref=True)
     assert dill.copy(Machine(), byref=True, recurse=True)
-    if not OLDER:
-        assert dill.copy(Machine(), recurse=True)
+    assert dill.copy(Machine(), recurse=True)
     assert dill.copy(Machine())
 
 
@@ -79,14 +78,28 @@ def __init__(self):
 def test_partials():
     assert dill.copy(SubMachine(), byref=True)
     assert dill.copy(SubMachine(), byref=True, recurse=True)
-   #if not OLDER: #FIXME: fails __main__.py
-   #    assert dill.copy(SubMachine(), recurse=True)
+    assert dill.copy(SubMachine(), recurse=True)
     assert dill.copy(SubMachine())
 
 
+class obj4(object):
+    def __init__(self):
+        super(obj4, self).__init__()
+        a = self
+        class obj5(object):
+            def __init__(self):
+                super(obj5, self).__init__()
+                self.a = a
+        self.b = obj5()
+
+
+def test_circular_reference():
+    assert dill.copy(obj4())
+
 
 if __name__ == '__main__':
     #print(('byref','_super','_recurse','_memo','_stop','OLDER'))
     test_super()
     test_partial()
     test_partials()
+    test_circular_reference()