Skip to content

Commit

Permalink
Merge pull request #72 from pitrou/py36
Browse files Browse the repository at this point in the history
Fix #71: make cloudpickle Python 3.6 compatible
  • Loading branch information
rgbkrk authored Nov 25, 2016
2 parents 304a0a4 + 319ece0 commit cbd3f34
Show file tree
Hide file tree
Showing 3 changed files with 100 additions and 41 deletions.
112 changes: 72 additions & 40 deletions cloudpickle/cloudpickle.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

Expand All @@ -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()
Expand Down Expand Up @@ -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

Expand Down
2 changes: 1 addition & 1 deletion tests/cloudpickle_file_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
27 changes: 27 additions & 0 deletions tests/cloudpickle_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
import pytest
import pickle
import sys
import random
import functools
import itertools
import platform
Expand Down Expand Up @@ -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()

0 comments on commit cbd3f34

Please sign in to comment.