support pypy

cloudpipe · May 25, 2017 · fb27c76 · fb27c76
1 parent 55cd6fd
commit fb27c76
Show file tree

Hide file tree

Showing 2 changed files with 101 additions and 61 deletions.
diff --git a/cloudpickle/cloudpickle.py b/cloudpickle/cloudpickle.py
@@ -56,6 +56,7 @@
 import types
 import weakref
 
+
 if sys.version < '3':
     from pickle import Pickler
     try:
@@ -70,50 +71,114 @@
     PY3 = True
 
 
-try:
-    from ctypes import pythonapi, py_object, c_int, PYFUNCTYPE
-except ImportError:
-    supports_recursive_closure = False
+def compress_closure(closure):
+    """Compress the closure by storing only the number of cells.
+    """
+    return len(closure) if closure is not None else -1
 
-    def compress_closure(closure):
-        return closure
 
-    def decompress_closure(compressed_closure):
-        return (
-            tuple(map(_make_cell, compressed_closure))
-            if compressed_closure is not None else
-            None
-        )
+def _make_cell_set_template_code():
+    """Get the Python compiler to emit LOAD_FAST(arg); STORE_DEREF
 
-    def save_closure(save, closure):
-        pass
+    Notes
+    -----
+    In Python 3, we could use an easier function:
 
-    def fill_cells(cells, values):
-        pass
-else:
-    supports_recursive_closure = True
+    .. code-block:: python
+
+       def f():
+           cell = None
+
+           def _stub(value):
+               nonlocal cell
+               cell = value
 
-    def compress_closure(closure):
-        return len(closure) if closure is not None else -1
+           return _stub
 
-    def decompress_closure(compressed_closure):
-        return (
-            tuple(_make_cell(None) for _ in range(compressed_closure))
-            if compressed_closure >= 0 else
-            None
+        _cell_set_template_code = f()
+
+    This function is _only_ a LOAD_FAST(arg); STORE_DEREF, but that is
+    invalid syntax on Python 2. If we use this function we also don't need
+    to do the weird freevars/cellvars swap below
+    """
+    def inner(value):
+        lambda: cell  # make ``cell`` a closure so that we get a STORE_DEREF
+        cell = value
+
+    co = inner.__code__
+
+    # NOTE: we are marking the cell variable as a free variable intentionally
+    # so that we simulate an inner function instead of the outer function. This
+    # is what gives us the ``nonlocal`` behavior in a Python 2 compatible way.
+    if not PY3:
+        return types.CodeType(
+            co.co_argcount,
+            co.co_nlocals,
+            co.co_stacksize,
+            co.co_flags,
+            co.co_code,
+            co.co_consts,
+            co.co_names,
+            co.co_varnames,
+            co.co_filename,
+            co.co_name,
+            co.co_firstlineno,
+            co.co_lnotab,
+            co.co_cellvars,
+            (),
+        )
+    else:
+        return types.CodeType(
+            co.co_argcount,
+            co.co_kwonlyargcount,
+            co.co_nlocals,
+            co.co_stacksize,
+            co.co_flags,
+            co.co_code,
+            co.co_consts,
+            co.co_names,
+            co.co_varnames,
+            co.co_filename,
+            co.co_name,
+            co.co_firstlineno,
+            co.co_lnotab,
+            co.co_cellvars,
+            (),
         )
 
-    def save_closure(save, closure):
-        save(closure)
 
-    _cell_set = PYFUNCTYPE(c_int, py_object, py_object)(
-        ('PyCell_Set', pythonapi), ((1, 'cell'), (1, 'value')),
-    )
+_cell_set_template_code = _make_cell_set_template_code()
+
+
+def _cell_set(cell, value):
+    """Set the value of a closure cell.
+    """
+    return types.FunctionType(
+        _cell_set_template_code,
+        {},
+        '_cell_set_inner',
+        (),
+        (cell,),
+    )(value)
+
+
+def fill_cells(cells, values):
+    """If we have a closure, fill the cells with their real values.
+    """
+    if cells is not None:
+        for cell, value in zip(cells, values):
+            _cell_set(cell, value)
+
 
-    def fill_cells(cells, values):
-        if cells is not None:
-            for cell, value in zip(cells, values):
-                _cell_set(cell, value)
+def decompress_closure(compressed_closure):
+    """Decompress the closure creating ``compressed_closure`` empty cells or
+    returning None.
+    """
+    return (
+        tuple(_make_cell(None) for _ in range(compressed_closure))
+        if compressed_closure >= 0 else
+        None
+    )
 
 
 #relevant opcodes
@@ -392,7 +457,7 @@ def save_function_tuple(self, func):
         save(f_globals)
         save(defaults)
         save(dct)
-        save_closure(save, closure)
+        save(closure)
         write(pickle.TUPLE)
         write(pickle.REDUCE)  # applies _fill_function on the tuple
 
@@ -869,10 +934,6 @@ def _make_cell(value):
     return (lambda: value).__closure__[0]
 
 
-def _reconstruct_closure(values):
-    return tuple([_make_cell(v) for v in values])
-
-
 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

diff --git a/tests/cloudpickle_test.py b/tests/cloudpickle_test.py
@@ -38,7 +38,7 @@
 from io import BytesIO
 
 import cloudpickle
-from cloudpickle.cloudpickle import _find_module, supports_recursive_closure
+from cloudpickle.cloudpickle import _find_module
 
 from .testutils import subprocess_pickle_echo
 
@@ -133,10 +133,6 @@ 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():
@@ -154,10 +150,6 @@ def g(n):
         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
@@ -175,19 +167,6 @@ def f():
             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"
-    )
-    def test_recursive_closure_unsupported(self):
-        def f1():
-            def g():
-                return g
-            return g
-
-        with pytest.raises(pickle.PicklingError):
-            pickle_depickle(f1())
-
     def test_unhashable_closure(self):
         def f():
             s = set((1, 2))  # mutable set is unhashable