recursive and unhashable closure

cloudpickle/cloudpickle.py

@@ -69,6 +69,43 @@
     from io import BytesIO as StringIO
     PY3 = True
 
+
+try:
+    from ctypes import pythonapi, py_object, c_int, PYFUNCTYPE
+except ImportError:
+    supports_recursive_closure = False
+
+    def compress_closure(closure):
+        return closure
+
+    def decompress_closure(compressed_closure):
+        return compressed_closure
+
+    def fill_cells(cells, values):
+        pass
+else:
+    supports_recursive_closure = True
+
+    def compress_closure(closure):
+        return len(closure) if closure is not None else -1
+
+    def decompress_closure(compressed_closure):
+        return (
+            tuple(_make_cell(None) for _ in range(compressed_closure))
+            if compressed_closure >= 0 else
+            None
+        )
+
+    _cell_set = PYFUNCTYPE(c_int, py_object, py_object)(
+        ('PyCell_Set', pythonapi), ((1, 'cell'), (1, 'value')),
+    )
+
+    def fill_cells(cells, values):
+        if cells is not None:
+            for cell, value in zip(cells, values):
+                _cell_set(cell, value)
+
+
 #relevant opcodes
 STORE_GLOBAL = opcode.opmap['STORE_GLOBAL']
 DELETE_GLOBAL = opcode.opmap['DELETE_GLOBAL']
@@ -305,7 +342,6 @@ def _save_subimports(self, code, top_level_dependencies):
                             # then discards the reference to it
                             self.write(pickle.POP)
 
-
     def save_function_tuple(self, func):
         """  Pickles an actual func object.
 
@@ -331,18 +367,22 @@ def save_function_tuple(self, func):
         save(_fill_function)  # skeleton function updater
         write(pickle.MARK)    # beginning of tuple that _fill_function expects
 
-        self._save_subimports(code, set(f_globals.values()) | set(closure))
+        self._save_subimports(
+            code,
+            itertools.chain(f_globals.values(), closure),
+        )
 
         # create a skeleton function object and memoize it
         save(_make_skel_func)
-        save((code, closure, base_globals))
+        save((code, compress_closure(func.__closure__), base_globals))
         write(pickle.REDUCE)
         self.memoize(func)
 
         # save the rest of the func data needed by _fill_function
         save(f_globals)
         save(defaults)
         save(dct)
+        save(closure)
         write(pickle.TUPLE)
         write(pickle.REDUCE)  # applies _fill_function on the tuple
 
@@ -799,14 +839,14 @@ def _gen_ellipsis():
 def _gen_not_implemented():
     return NotImplemented
 
-def _fill_function(func, globals, defaults, dict):
+def _fill_function(func, globals, defaults, dict, closure):
     """ Fills in the rest of function data into the skeleton function object
         that were created via _make_skel_func().
          """
     func.__globals__.update(globals)
     func.__defaults__ = defaults
     func.__dict__ = dict
-
+    fill_cells(func.__closure__, closure)
     return func
 
 
@@ -818,19 +858,17 @@ def _reconstruct_closure(values):
     return tuple([_make_cell(v) for v in values])
 
 
-def _make_skel_func(code, closures, base_globals = None):
+def _make_skel_func(code, compressed_closure, base_globals=None):
     """ Creates a skeleton function object that contains just the provided
         code and the correct number of cells in func_closure.  All other
         func attributes (e.g. func_globals) are empty.
     """
-    closure = _reconstruct_closure(closures) if closures else None
-
     if base_globals is None:
         base_globals = {}
     base_globals['__builtins__'] = __builtins__
 
-    return types.FunctionType(code, base_globals,
-                              None, None, closure)
+    closure = decompress_closure(compressed_closure)
+    return types.FunctionType(code, base_globals, None, None, closure)
 
 
 def _find_module(mod_name):

tests/cloudpickle_test.py

@@ -38,7 +38,7 @@
 from io import BytesIO
 
 import cloudpickle
-from cloudpickle.cloudpickle import _find_module
+from cloudpickle.cloudpickle import _find_module, supports_recursive_closure
 
 from .testutils import subprocess_pickle_echo
 
@@ -133,6 +133,70 @@ def test_nested_lambdas(self):
         f2 = lambda x: f1(x) // b
         self.assertEqual(pickle_depickle(f2)(1), 1)
 
+    @pytest.mark.skipif(
+        not supports_recursive_closure,
+        reason='The C API is needed for recursively defined closures'
+    )
+    def test_recursive_closure(self):
+        def f1():
+            def g():
+                return g
+            return g
+
+        def f2(base):
+            def g(n):
+                return base if n <= 1 else n * g(n - 1)
+            return g
+
+        g1 = pickle_depickle(f1())
+        self.assertEqual(g1(), g1)
+
+        g2 = pickle_depickle(f2(2))
+        self.assertEqual(g2(5), 240)
+
+    @pytest.mark.skipif(
+        not supports_recursive_closure,
+        reason='The C API is needed for recursively defined closures'
+    )
+    def test_closure_none_is_preserved(self):
+        def f():
+            """a function with no closure cells
+            """
+
+        self.assertIsNone(f.__closure__, msg='f actually has closure cells!')
+
+        g = pickle_depickle(f)
+
+        self.assertIsNone(
+            g.__closure__,
+            msg='g now has closure cells even though f does not',
+        )
+
+    @pytest.mark.skipif(
+        supports_recursive_closure,
+        reason="Recursive closures shouldn't raise an exception if supported"
+    )
+    @pytest.mark.xfail
+    def test_recursive_closure_unsupported(self):
+        def f1():
+            def g():
+                return g
+            return g
+
+        pickle_depickle(f1())
+
+    def test_unhashable_closure(self):
+        def f():
+            s = set((1, 2))  # mutable set is unhashable
+
+            def g():
+                return len(s)
+
+            return g
+
+        g = pickle_depickle(f())
+        self.assertEqual(g(), 2)
+
     @pytest.mark.skipif(sys.version_info >= (3, 4)
                         and sys.version_info < (3, 4, 3),
                         reason="subprocess has a bug in 3.4.0 to 3.4.2")