| """ |
| This class is defined to override standard pickle functionality |
| The goals of it follow: |
| -Serialize lambdas and nested functions to compiled byte code |
| -Deal with main module correctly |
| -Deal with other non-serializable objects |
| It does not include an unpickler, as standard python unpickling suffices. |
| This module was extracted from the `cloud` package, developed by `PiCloud, Inc. |
| <http://www.picloud.com>`_. |
| Copyright (c) 2012, Regents of the University of California. |
| Copyright (c) 2009 `PiCloud, Inc. <http://www.picloud.com>`_. |
| All rights reserved. |
| Redistribution and use in source and binary forms, with or without |
| modification, are permitted provided that the following conditions |
| are met: |
| * Redistributions of source code must retain the above copyright |
| notice, this list of conditions and the following disclaimer. |
| * Redistributions in binary form must reproduce the above copyright |
| notice, this list of conditions and the following disclaimer in the |
| documentation and/or other materials provided with the distribution. |
| * Neither the name of the University of California, Berkeley nor the |
| names of its contributors may be used to endorse or promote |
| products derived from this software without specific prior written |
| permission. |
| THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS |
| "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT |
| LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR |
| A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT |
| HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, |
| SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED |
| TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR |
| PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF |
| LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING |
| NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS |
| SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. |
| """ |
| from __future__ import print_function |
| |
| import operator |
| import opcode |
| import os |
| import io |
| import pickle |
| import struct |
| import sys |
| import types |
| from functools import partial |
| import itertools |
| import dis |
| import traceback |
| import weakref |
| |
| # pylint: disable-all |
| |
| if sys.version < '3': |
| from pickle import Pickler # pylint: disable=ungrouped-imports |
| try: |
| from cStringIO import StringIO |
| except ImportError: |
| from StringIO import StringIO |
| PY3 = False |
| else: |
| types.ClassType = type |
| from pickle import _Pickler as Pickler # pylint: disable=ungrouped-imports |
| from io import BytesIO as StringIO # pylint: disable=ungrouped-imports |
| PY3 = True |
| |
| #relevant opcodes |
| 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 |
| |
| |
| def islambda(func): |
| return getattr(func, '__name__') == '<lambda>' |
| |
| |
| _BUILTIN_TYPE_NAMES = {} |
| for k1, v1 in types.__dict__.items(): |
| if type(v1) is type: # pylint: disable=unidiomatic-typecheck |
| _BUILTIN_TYPE_NAMES[v1] = k1 |
| |
| |
| 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): # pylint: disable=no-member |
| op = instr.opcode |
| if op in GLOBAL_OPS: |
| yield op, instr.arg |
| |
| |
| class CloudPickler(Pickler): # pylint: disable=too-many-public-methods |
| """ |
| CloudPickler class |
| """ |
| dispatch = Pickler.dispatch.copy() |
| |
| def __init__(self, filen, protocol=None): |
| Pickler.__init__(self, filen, protocol) |
| # set of modules to unpickle |
| self.modules = set() |
| # map ids to dictionary. used to ensure that functions can share global env |
| self.globals_ref = {} |
| |
| def dump(self, obj): |
| self.inject_addons() |
| try: |
| return Pickler.dump(self, obj) |
| except RuntimeError as e: |
| if 'recursion' in e.args[0]: |
| msg = """Could not pickle object as excessively deep recursion required.""" |
| raise pickle.PicklingError(msg) |
| except pickle.PickleError: |
| raise |
| except Exception as e: |
| print_exec(sys.stderr) |
| raise pickle.PicklingError(str(e)) |
| |
| def save_memoryview(self, obj): |
| """Fallback to save_string""" |
| Pickler.save_string(self, str(obj)) |
| |
| def save_buffer(self, obj): |
| """Fallback to save_string""" |
| Pickler.save_string(self, str(obj)) |
| if PY3: |
| dispatch[memoryview] = save_memoryview |
| else: |
| dispatch[buffer] = save_buffer |
| |
| def save_unsupported(self, obj): # pylint: disable=no-self-use |
| raise pickle.PicklingError("Cannot pickle objects of type %s" % type(obj)) |
| dispatch[types.GeneratorType] = save_unsupported |
| |
| # itertools objects do not pickle! |
| for v in itertools.__dict__.values(): |
| if type(v) is type: # pylint: disable=unidiomatic-typecheck |
| dispatch[v] = save_unsupported |
| |
| def save_module(self, obj): |
| """ |
| Save a module as an import |
| """ |
| self.modules.add(obj) |
| self.save_reduce(subimport, (obj.__name__,), obj=obj) |
| dispatch[types.ModuleType] = save_module |
| |
| def save_codeobject(self, obj): |
| """ |
| Save a code object |
| """ |
| if PY3: |
| args = ( |
| obj.co_argcount, obj.co_kwonlyargcount, obj.co_nlocals, obj.co_stacksize, |
| obj.co_flags, obj.co_code, obj.co_consts, obj.co_names, obj.co_varnames, |
| obj.co_filename, obj.co_name, obj.co_firstlineno, obj.co_lnotab, obj.co_freevars, |
| obj.co_cellvars |
| ) |
| else: |
| args = ( |
| obj.co_argcount, obj.co_nlocals, obj.co_stacksize, obj.co_flags, obj.co_code, |
| obj.co_consts, obj.co_names, obj.co_varnames, obj.co_filename, obj.co_name, |
| obj.co_firstlineno, obj.co_lnotab, obj.co_freevars, obj.co_cellvars |
| ) |
| self.save_reduce(types.CodeType, args, obj=obj) |
| dispatch[types.CodeType] = save_codeobject |
| |
| def save_function(self, obj, name=None): |
| """ Registered with the dispatch to handle all function types. |
| Determines what kind of function obj is (e.g. lambda, defined at |
| interactive prompt, etc) and handles the pickling appropriately. |
| """ |
| write = self.write |
| |
| if name is None: |
| name = obj.__name__ |
| try: |
| # whichmodule() could fail, see |
| # https://bitbucket.org/gutworth/six/issues/63/importing-six-breaks-pickling |
| modname = pickle.whichmodule(obj, name) |
| except Exception: |
| modname = None |
| # print('which gives %s %s %s' % (modname, obj, name)) |
| try: |
| themodule = sys.modules[modname] |
| except KeyError: |
| # eval'd items such as namedtuple give invalid items for their function __module__ |
| modname = '__main__' |
| |
| if modname == '__main__': |
| themodule = None |
| |
| if themodule: |
| self.modules.add(themodule) |
| if getattr(themodule, name, None) is obj: |
| return self.save_global(obj, name) |
| |
| # if func is lambda, def'ed at prompt, is in main, or is nested, then |
| # we'll pickle the actual function object rather than simply saving a |
| # reference (as is done in default pickler), via save_function_tuple. |
| if islambda(obj) or obj.__code__.co_filename == '<stdin>' or themodule is None: |
| #print("save global", islambda(obj), obj.__code__.co_filename, modname, themodule) |
| self.save_function_tuple(obj) |
| return |
| else: |
| # func is nested |
| klass = getattr(themodule, name, None) |
| if klass is None or klass is not obj: |
| self.save_function_tuple(obj) |
| return |
| |
| if obj.__dict__: |
| # essentially save_reduce, but workaround needed to avoid recursion |
| self.save(_restore_attr) |
| write(pickle.MARK + pickle.GLOBAL + modname + '\n' + name + '\n') |
| self.memoize(obj) |
| self.save(obj.__dict__) |
| write(pickle.TUPLE + pickle.REDUCE) |
| else: |
| write(pickle.GLOBAL + modname + '\n' + name + '\n') |
| self.memoize(obj) |
| dispatch[types.FunctionType] = save_function |
| |
| def save_function_tuple(self, func): |
| """ Pickles an actual func object. |
| A func comprises: code, globals, defaults, closure, and dict. We |
| extract and save these, injecting reducing functions at certain points |
| to recreate the func object. Keep in mind that some of these pieces |
| can contain a ref to the func itself. Thus, a naive save on these |
| pieces could trigger an infinite loop of save's. To get around that, |
| we first create a skeleton func object using just the code (this is |
| safe, since this won't contain a ref to the func), and memoize it as |
| soon as it's created. The other stuff can then be filled in later. |
| """ |
| save = self.save |
| write = self.write |
| |
| code, f_globals, defaults, closure, dct, base_globals = self.extract_func_data(func) |
| |
| save(_fill_function) # skeleton function updater |
| write(pickle.MARK) # beginning of tuple that _fill_function expects |
| |
| # create a skeleton function object and memoize it |
| save(_make_skel_func) |
| save((code, 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(func.__module__) |
| write(pickle.TUPLE) |
| write(pickle.REDUCE) # applies _fill_function on the tuple |
| |
| _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)) |
| |
| # see if nested function have any global refs |
| if co.co_consts: |
| for const in co.co_consts: |
| if type(const) is types.CodeType: # pylint: disable=unidiomatic-typecheck |
| out_names |= cls.extract_code_globals(const) |
| |
| cls._extract_code_globals_cache[co] = out_names |
| |
| return out_names |
| |
| def extract_func_data(self, func): |
| """ |
| Turn the function into a tuple of data necessary to recreate it: |
| code, globals, defaults, closure, dict |
| """ |
| code = func.__code__ |
| |
| # extract all global ref's |
| func_global_refs = self.extract_code_globals(code) |
| |
| # process all variables referenced by global environment |
| f_globals = {} |
| for var in func_global_refs: |
| if var in func.__globals__: |
| f_globals[var] = func.__globals__[var] |
| |
| # defaults requires no processing |
| defaults = func.__defaults__ |
| |
| # process closure |
| closure = [c.cell_contents for c in func.__closure__] if func.__closure__ else [] |
| |
| # save the dict |
| dct = func.__dict__ |
| |
| base_globals = self.globals_ref.get(id(func.__globals__), {}) |
| self.globals_ref[id(func.__globals__)] = base_globals |
| |
| return (code, f_globals, defaults, closure, dct, base_globals) |
| |
| def save_builtin_function(self, obj): |
| if obj.__module__ is "__builtin__": |
| return self.save_global(obj) |
| return self.save_function(obj) |
| dispatch[types.BuiltinFunctionType] = save_builtin_function |
| |
| def save_global(self, obj, name=None, pack=struct.pack): # pylint: disable=too-many-branches |
| if obj.__module__ == "__builtin__" or obj.__module__ == "builtins": |
| if obj in _BUILTIN_TYPE_NAMES: |
| return self.save_reduce(_builtin_type, (_BUILTIN_TYPE_NAMES[obj],), obj=obj) |
| |
| if name is None: |
| name = obj.__name__ |
| |
| modname = getattr(obj, "__module__", None) |
| if modname is None: |
| try: |
| # whichmodule() could fail, see |
| # https://bitbucket.org/gutworth/six/issues/63/importing-six-breaks-pickling |
| modname = pickle.whichmodule(obj, name) |
| except Exception: |
| modname = '__main__' |
| |
| if modname == '__main__': |
| themodule = None |
| else: |
| __import__(modname) |
| themodule = sys.modules[modname] |
| self.modules.add(themodule) |
| |
| if hasattr(themodule, name) and getattr(themodule, name) is obj: |
| return Pickler.save_global(self, obj, name) |
| |
| typ = type(obj) |
| if typ is not obj and isinstance(obj, (type, types.ClassType)): |
| d = dict(obj.__dict__) # copy dict proxy to a dict |
| if not isinstance(d.get('__dict__', None), property): |
| # don't extract dict that are properties |
| d.pop('__dict__', None) |
| d.pop('__weakref__', None) |
| |
| # hack as __new__ is stored differently in the __dict__ |
| new_override = d.get('__new__', None) |
| if new_override: |
| d['__new__'] = obj.__new__ |
| |
| # workaround for namedtuple (hijacked by PySpark) |
| if getattr(obj, '_is_namedtuple_', False): |
| self.save_reduce(_load_namedtuple, (obj.__name__, obj._fields)) |
| return |
| |
| self.save(_load_class) |
| self.save_reduce(typ, (obj.__name__, obj.__bases__, {"__doc__": obj.__doc__}), obj=obj) |
| d.pop('__doc__', None) |
| # handle property and staticmethod |
| dd = {} |
| for k, v in d.items(): |
| if isinstance(v, property): |
| k = ('property', k) |
| v = (v.fget, v.fset, v.fdel, v.__doc__) |
| elif isinstance(v, staticmethod) and hasattr(v, '__func__'): |
| k = ('staticmethod', k) |
| v = v.__func__ |
| elif isinstance(v, classmethod) and hasattr(v, '__func__'): |
| k = ('classmethod', k) |
| v = v.__func__ |
| dd[k] = v |
| self.save(dd) |
| self.write(pickle.TUPLE2) |
| self.write(pickle.REDUCE) |
| |
| else: |
| raise pickle.PicklingError("Can't pickle %r" % obj) |
| |
| dispatch[type] = save_global |
| dispatch[types.ClassType] = save_global |
| |
| def save_instancemethod(self, obj): |
| # Memoization rarely is ever useful due to python bounding |
| if PY3: |
| self.save_reduce(types.MethodType, (obj.__func__, obj.__self__), obj=obj) |
| else: |
| self.save_reduce( |
| types.MethodType, (obj.__func__, obj.__self__, obj.__self__.__class__), |
| obj=obj) |
| dispatch[types.MethodType] = save_instancemethod |
| |
| def save_inst(self, obj): |
| """Inner logic to save instance. Based off pickle.save_inst |
| Supports __transient__""" |
| cls = obj.__class__ |
| |
| memo = self.memo |
| write = self.write |
| save = self.save |
| |
| if hasattr(obj, '__getinitargs__'): |
| args = obj.__getinitargs__() |
| len(args) # assert it's a sequence |
| pickle._keep_alive(args, memo) # pylint: disable=protected-access |
| else: |
| args = () |
| |
| write(pickle.MARK) |
| |
| if self.bin: |
| save(cls) |
| for arg in args: |
| save(arg) |
| write(pickle.OBJ) |
| else: |
| for arg in args: |
| save(arg) |
| write(pickle.INST + cls.__module__ + '\n' + cls.__name__ + '\n') |
| |
| self.memoize(obj) |
| |
| try: |
| getstate = obj.__getstate__ |
| except AttributeError: |
| stuff = obj.__dict__ |
| #remove items if transient |
| if hasattr(obj, '__transient__'): |
| transient = obj.__transient__ |
| stuff = stuff.copy() |
| for k in list(stuff.keys()): |
| if k in transient: |
| del stuff[k] |
| else: |
| stuff = getstate() |
| pickle._keep_alive(stuff, memo) # pylint: disable=protected-access |
| save(stuff) |
| write(pickle.BUILD) |
| |
| if not PY3: |
| dispatch[types.InstanceType] = save_inst |
| |
| def save_property(self, obj): |
| # properties not correctly saved in python |
| self.save_reduce(property, (obj.fget, obj.fset, obj.fdel, obj.__doc__), obj=obj) |
| dispatch[property] = save_property |
| |
| def save_itemgetter(self, obj): |
| """itemgetter serializer (needed for namedtuple support)""" |
| class Dummy: # pylint: disable=old-style-class |
| def __init__(self): |
| pass |
| def __getitem__(self, item): |
| return item |
| items = obj(Dummy()) |
| if not isinstance(items, tuple): |
| items = (items, ) |
| return self.save_reduce(operator.itemgetter, items) |
| |
| if type(operator.itemgetter) is type: # pylint: disable=unidiomatic-typecheck |
| dispatch[operator.itemgetter] = save_itemgetter |
| |
| def save_attrgetter(self, obj): |
| """attrgetter serializer""" |
| class Dummy(object): |
| def __init__(self, attrs, index=None): |
| self.attrs = attrs |
| self.index = index |
| def __getattribute__(self, item): |
| attrs = object.__getattribute__(self, "attrs") |
| index = object.__getattribute__(self, "index") |
| if index is None: |
| index = len(attrs) |
| attrs.append(item) |
| else: |
| attrs[index] = ".".join([attrs[index], item]) |
| return type(self)(attrs, index) |
| attrs = [] |
| obj(Dummy(attrs)) |
| return self.save_reduce(operator.attrgetter, tuple(attrs)) |
| |
| if type(operator.attrgetter) is type: # pylint: disable=unidiomatic-typecheck |
| dispatch[operator.attrgetter] = save_attrgetter |
| |
| def save_reduce(self, func, args, state=None, # pylint: disable=too-many-branches |
| listitems=None, dictitems=None, obj=None): |
| """Modified to support __transient__ on new objects |
| Change only affects protocol level 2 (which is always used by PiCloud""" |
| # Assert that args is a tuple or None |
| if not isinstance(args, tuple): |
| raise pickle.PicklingError("args from reduce() should be a tuple") |
| |
| # Assert that func is callable |
| if not hasattr(func, '__call__'): |
| raise pickle.PicklingError("func from reduce should be callable") |
| |
| save = self.save |
| write = self.write |
| |
| # Protocol 2 special case: if func's name is __newobj__, use NEWOBJ |
| if self.proto >= 2 and getattr(func, "__name__", "") == "__newobj__": |
| #Added fix to allow transient |
| cls = args[0] |
| if not hasattr(cls, "__new__"): |
| raise pickle.PicklingError( |
| "args[0] from __newobj__ args has no __new__") |
| if obj is not None and cls is not obj.__class__: |
| raise pickle.PicklingError( |
| "args[0] from __newobj__ args has the wrong class") |
| args = args[1:] |
| save(cls) |
| |
| #Don't pickle transient entries |
| if hasattr(obj, '__transient__'): |
| transient = obj.__transient__ |
| state = state.copy() |
| |
| for k in list(state.keys()): |
| if k in transient: |
| del state[k] |
| |
| save(args) |
| write(pickle.NEWOBJ) |
| else: |
| save(func) |
| save(args) |
| write(pickle.REDUCE) |
| |
| if obj is not None: |
| self.memoize(obj) |
| |
| # More new special cases (that work with older protocols as |
| # well): when __reduce__ returns a tuple with 4 or 5 items, |
| # the 4th and 5th item should be iterators that provide list |
| # items and dict items (as (key, value) tuples), or None. |
| |
| if listitems is not None: |
| self._batch_appends(listitems) |
| |
| if dictitems is not None: |
| self._batch_setitems(dictitems) |
| |
| if state is not None: |
| save(state) |
| write(pickle.BUILD) |
| |
| def save_partial(self, obj): |
| """Partial objects do not serialize correctly in python2.x -- this fixes the bugs""" |
| self.save_reduce(_genpartial, (obj.func, obj.args, obj.keywords)) |
| |
| if sys.version_info < (2, 7): # 2.7 supports partial pickling |
| dispatch[partial] = save_partial |
| |
| |
| def save_file(self, obj): # pylint: disable=too-many-branches |
| """Save a file""" |
| try: |
| import StringIO as pystringIO #we can't use cStringIO as it lacks the name attribute |
| except ImportError: |
| import io as pystringIO # pylint: disable=reimported |
| |
| if not hasattr(obj, 'name') or not hasattr(obj, 'mode'): |
| raise pickle.PicklingError("Cannot pickle files that do not map to an actual file") |
| if obj is sys.stdout: |
| return self.save_reduce(getattr, (sys, 'stdout'), obj=obj) |
| if obj is sys.stderr: |
| return self.save_reduce(getattr, (sys, 'stderr'), obj=obj) |
| if obj is sys.stdin: |
| raise pickle.PicklingError("Cannot pickle standard input") |
| if hasattr(obj, 'isatty') and obj.isatty(): |
| raise pickle.PicklingError("Cannot pickle files that map to tty objects") |
| if 'r' not in obj.mode: |
| raise pickle.PicklingError("Cannot pickle files that are not opened for reading") |
| name = obj.name |
| try: |
| fsize = os.stat(name).st_size |
| except OSError: |
| raise pickle.PicklingError("Cannot pickle file %s as it cannot be stat" % name) |
| |
| if obj.closed: |
| #create an empty closed string io |
| retval = pystringIO.StringIO("") |
| retval.close() |
| elif not fsize: #empty file |
| retval = pystringIO.StringIO("") |
| try: |
| tmpfile = file(name) |
| tst = tmpfile.read(1) |
| except IOError: |
| raise pickle.PicklingError("Cannot pickle file %s as it cannot be read" % name) |
| tmpfile.close() |
| if tst != '': |
| raise pickle.PicklingError( |
| "Cannot pickle file %s as it does not appear to map to a physical, real file" % name) |
| else: |
| try: |
| tmpfile = file(name) |
| contents = tmpfile.read() |
| tmpfile.close() |
| except IOError: |
| raise pickle.PicklingError("Cannot pickle file %s as it cannot be read" % name) |
| retval = pystringIO.StringIO(contents) |
| curloc = obj.tell() |
| retval.seek(curloc) |
| |
| retval.name = name |
| self.save(retval) |
| self.memoize(obj) |
| |
| if PY3: |
| dispatch[io.TextIOWrapper] = save_file |
| else: |
| dispatch[file] = save_file |
| |
| # Special functions for Add-on libraries |
| |
| def inject_numpy(self): |
| numpy = sys.modules.get('numpy') |
| if not numpy or not hasattr(numpy, 'ufunc'): |
| return |
| self.dispatch[numpy.ufunc] = self.__class__.save_ufunc |
| |
| def save_ufunc(self, obj): |
| """Hack function for saving numpy ufunc objects""" |
| name = obj.__name__ |
| numpy_tst_mods = ['numpy', 'scipy.special'] |
| for tst_mod_name in numpy_tst_mods: |
| tst_mod = sys.modules.get(tst_mod_name, None) |
| if tst_mod and name in tst_mod.__dict__: |
| return self.save_reduce(_getobject, (tst_mod_name, name)) |
| raise pickle.PicklingError( |
| 'cannot save %s. Cannot resolve what module it is defined in' % str(obj)) |
| |
| def inject_addons(self): |
| """Plug in system. Register additional pickling functions if modules already loaded""" |
| self.inject_numpy() |
| |
| |
| # Shorthands for legacy support |
| |
| def dump(obj, filen, protocol=2): |
| CloudPickler(filen, protocol).dump(obj) |
| |
| |
| def dumps(obj, protocol=2): |
| filen = StringIO() |
| |
| cp = CloudPickler(filen, protocol) |
| cp.dump(obj) |
| |
| return filen.getvalue() |
| |
| |
| #hack for __import__ not working as desired |
| def subimport(name): |
| __import__(name) |
| return sys.modules[name] |
| |
| |
| # restores function attributes |
| def _restore_attr(obj, attr): |
| for key, val in attr.items(): |
| setattr(obj, key, val) |
| return obj |
| |
| |
| def _get_module_builtins(): |
| return pickle.__builtins__ # pylint: disable=no-member |
| |
| |
| def print_exec(stream): |
| ei = sys.exc_info() |
| traceback.print_exception(ei[0], ei[1], ei[2], None, stream) |
| |
| |
| def _modules_to_main(modList): |
| """Force every module in modList to be placed into main""" |
| if not modList: |
| return |
| |
| main = sys.modules['__main__'] |
| for modname in modList: |
| if isinstance(modname, str): |
| try: |
| mod = __import__(modname) |
| except Exception: |
| sys.stderr.write( |
| 'warning: could not import %s\n. ' |
| 'Your function may unexpectedly error due to this import failing;' |
| 'A version mismatch is likely. Specific error was:\n' % modname) |
| print_exec(sys.stderr) |
| else: |
| setattr(main, mod.__name__, mod) |
| |
| |
| #object generators: |
| def _genpartial(func, args, kwds): |
| if not args: |
| args = () |
| if not kwds: |
| kwds = {} |
| return partial(func, *args, **kwds) |
| |
| |
| def _fill_function(func, globalsn, defaults, dictn, module): |
| """ Fills in the rest of function data into the skeleton function object |
| that were created via _make_skel_func(). |
| """ |
| func.__globals__.update(globalsn) |
| func.__defaults__ = defaults |
| func.__dict__ = dictn |
| func.__module__ = module |
| |
| return func |
| |
| |
| 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, closures, 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) |
| |
| |
| def _load_class(cls, d): |
| """ |
| Loads additional properties into class `cls`. |
| """ |
| for k, v in d.items(): |
| if isinstance(k, tuple): |
| typ, k = k |
| if typ == 'property': |
| v = property(*v) |
| elif typ == 'staticmethod': |
| v = staticmethod(v) # pylint: disable=redefined-variable-type |
| elif typ == 'classmethod': |
| v = classmethod(v) |
| setattr(cls, k, v) |
| return cls |
| |
| |
| def _load_namedtuple(name, fields): |
| """ |
| Loads a class generated by namedtuple |
| """ |
| from collections import namedtuple |
| return namedtuple(name, fields) |
| |
| |
| # Constructors for 3rd party libraries |
| # Note: These can never be renamed due to client compatibility issues |
| |
| def _getobject(modname, attribute): |
| mod = __import__(modname, fromlist=[attribute]) |
| return mod.__dict__[attribute] |