Fix #71: make cloudpickle Python 3.6 compatible

cloudpickle/cloudpickle.py

@@ -42,17 +42,19 @@
 """
 from __future__ import print_function
 
-import operator
-import io
+import dis
+from functools import partial
 import imp
+import io
+import itertools
+import opcode
+import operator
 import pickle
 import struct
 import sys
-import types
-from functools import partial
-import itertools
-import dis
 import traceback
+import types
+import weakref
 
 if sys.version < '3':
     from pickle import Pickler
@@ -68,10 +70,10 @@
     PY3 = True
 
 #relevant opcodes
-STORE_GLOBAL = dis.opname.index('STORE_GLOBAL')
-DELETE_GLOBAL = dis.opname.index('DELETE_GLOBAL')
-LOAD_GLOBAL = dis.opname.index('LOAD_GLOBAL')
-GLOBAL_OPS = [STORE_GLOBAL, DELETE_GLOBAL, LOAD_GLOBAL]
+STORE_GLOBAL = opcode.opmap['STORE_GLOBAL']
+DELETE_GLOBAL = opcode.opmap['DELETE_GLOBAL']
+LOAD_GLOBAL = opcode.opmap['LOAD_GLOBAL']
+GLOBAL_OPS = (STORE_GLOBAL, DELETE_GLOBAL, LOAD_GLOBAL)
 HAVE_ARGUMENT = dis.HAVE_ARGUMENT
 EXTENDED_ARG = dis.EXTENDED_ARG
 
@@ -90,6 +92,43 @@ def _builtin_type(name):
     return getattr(types, name)
 
 
+if sys.version_info < (3, 4):
+    def _walk_global_ops(code):
+        """
+        Yield (opcode, argument number) tuples for all
+        global-referencing instructions in *code*.
+        """
+        code = getattr(code, 'co_code', b'')
+        if not PY3:
+            code = map(ord, code)
+
+        n = len(code)
+        i = 0
+        extended_arg = 0
+        while i < n:
+            op = code[i]
+            i += 1
+            if op >= HAVE_ARGUMENT:
+                oparg = code[i] + code[i + 1] * 256 + extended_arg
+                extended_arg = 0
+                i += 2
+                if op == EXTENDED_ARG:
+                    extended_arg = oparg * 65536
+                if op in GLOBAL_OPS:
+                    yield op, oparg
+
+else:
+    def _walk_global_ops(code):
+        """
+        Yield (opcode, argument number) tuples for all
+        global-referencing instructions in *code*.
+        """
+        for instr in dis.get_instructions(code):
+            op = instr.opcode
+            if op in GLOBAL_OPS:
+                yield op, instr.arg
+
+
 class CloudPickler(Pickler):
 
     dispatch = Pickler.dispatch.copy()
@@ -281,41 +320,34 @@ def save_function_tuple(self, func):
         write(pickle.TUPLE)
         write(pickle.REDUCE)  # applies _fill_function on the tuple
 
-    @staticmethod
-    def extract_code_globals(co):
+    _extract_code_globals_cache = (
+        weakref.WeakKeyDictionary()
+        if sys.version_info >= (2, 7) and not hasattr(sys, "pypy_version_info")
+        else {})
+
+    @classmethod
+    def extract_code_globals(cls, co):
         """
         Find all globals names read or written to by codeblock co
         """
+        out_names = cls._extract_code_globals_cache.get(co)
+        if out_names is None:
+            try:
+                names = co.co_names
+            except AttributeError:
+                # PyPy "builtin-code" object
+                out_names = set()
+            else:
+                out_names = set(names[oparg]
+                                for op, oparg in _walk_global_ops(co))
 
-        code = getattr(co, 'co_code', None)
-        if code is None:
-            return set()
-        if not PY3:
-            code = [ord(c) for c in code]
-        names = co.co_names
-        out_names = set()
-
-        n = len(code)
-        i = 0
-        extended_arg = 0
-        while i < n:
-            op = code[i]
-
-            i += 1
-            if op >= HAVE_ARGUMENT:
-                oparg = code[i] + code[i+1] * 256 + extended_arg
-                extended_arg = 0
-                i += 2
-                if op == EXTENDED_ARG:
-                    extended_arg = oparg*65536
-                if op in GLOBAL_OPS:
-                    out_names.add(names[oparg])
+                # see if nested function have any global refs
+                if co.co_consts:
+                    for const in co.co_consts:
+                        if type(const) is types.CodeType:
+                            out_names |= cls.extract_code_globals(const)
 
-        # see if nested function have any global refs
-        if co.co_consts:
-            for const in co.co_consts:
-                if type(const) is types.CodeType:
-                    out_names |= CloudPickler.extract_code_globals(const)
+            cls._extract_code_globals_cache[co] = out_names
 
         return out_names
 

tests/cloudpickle_file_test.py

@@ -85,7 +85,7 @@ def test_seek(self):
             self.assertEquals(self.teststring, unpickled.read())
         os.remove(self.tmpfilepath)
 
-    @pytest.mark.skipif(sys.version_info > (2, 7),
+    @pytest.mark.skipif(sys.version_info >= (3,),
                         reason="only works on Python 2.x")
     def test_temp_file(self):
         with tempfile.NamedTemporaryFile(mode='ab+') as fp:

tests/cloudpickle_test.py

@@ -4,6 +4,7 @@
 import pytest
 import pickle
 import sys
+import random
 import functools
 import itertools
 import platform
@@ -333,6 +334,32 @@ def g(y):
         res = loop.run_sync(functools.partial(g2, 5))
         self.assertEqual(res, 7)
 
+    def test_extended_arg(self):
+        # Functions with more than 65535 global vars prefix some global
+        # variable references with the EXTENDED_ARG opcode.
+        nvars = 65537 + 258
+        names = ['g%d' % i for i in range(1, nvars)]
+        r = random.Random(42)
+        d = dict([(name, r.randrange(100)) for name in names])
+        # def f(x):
+        #     x = g1, g2, ...
+        #     return zlib.crc32(bytes(bytearray(x)))
+        code = """
+        import zlib
+
+        def f():
+            x = {tup}
+            return zlib.crc32(bytes(bytearray(x)))
+        """.format(tup=', '.join(names))
+        exec(textwrap.dedent(code), d, d)
+        f = d['f']
+        res = f()
+        data = cloudpickle.dumps([f, f])
+        d = f = None
+        f2, f3 = pickle.loads(data)
+        self.assertTrue(f2 is f3)
+        self.assertEqual(f2(), res)
+
 
 if __name__ == '__main__':
     unittest.main()