Source code for brine

# This library is free software; you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as
# published by the Free Software Foundation; either version 3 of the
# License, or (at your option) any later version.
#
# This library is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public
# License along with this library; if not, see
# <http://www.gnu.org/licenses/>.


"""
Brine provides a way to wrap function objects so that they may be
pickled.

To truly pickle a function we need to be able to duplicate its code
and its closures. By default, `pickle` will simply store the
function's name, and then attempt to associate that with a function
when unpickling. This of course fails when the function is a lambda or
not otherwise defined at the top level.

In order to mark a function, method, or partial for storage, use the
`brine` function to create a wrapper. Later, after pickling and
unpickling the wrapper, call `unbrine` to get a new copy of the
original function.

Loading this module has the side effect of registering a pickle
handler for the `CellType` and `CodeType` types. This should be of low
impact, as the only place these types are used is within function
instances, and are typically unexposed.

:author: Christopher O'Brien  <obriencj@gmail.com>
:license: LGPL v.3
"""


from abc import ABCMeta, abstractmethod
from functools import partial
from types import BuiltinFunctionType, BuiltinMethodType
from types import FunctionType, MethodType, CodeType
from ._cellwork import CellType, cell_get_value, cell_from_value

import copy_reg
import new


__all__ = [ "brine", "unbrine",
            "BrinedObject",
            "BrinedFunction", "BrinedMethod", "BrinedPartial",
            "code_unnew", "code_new",
            "function_unnew", "function_new" ]


[docs]def brine(value): """ Wrap an object so that it may be pickled. Behavior by type of `value` is as follows: * `builtin_function_or_method` is unchanged * `partial` is wrapped as a `BrinedPartial` * `instancemethod` is wrapped as a `BrinedMethod` * `function` is wrapped as a `BrinedFunction` * `list` and `tuple` are duplicated and contents are brined * `dict` is duplicated and its values are brined * all other types are returned unchanged This function provides neither caching nor preservation of uniqueness -- if the same function is in a list multiple times, each will be wrapped individually and as a result will be duplicated when unbrined. In cases where uniqueness needs to be preserved, use a `Barrel` instead. Also, if a brined function refers to other functions in its closure cells, they will not be brined. Use a `Barrel` for those cases as well. Parameters ---------- value : `object` object to be brined for pickling Returns ------- wrapped : `object` as defined by the `value` parameter """ if isinstance(value, (BuiltinFunctionType, BuiltinMethodType)): return value elif isinstance(value, partial): return BrinedPartial(value) elif isinstance(value, MethodType): return BrinedMethod(value) elif isinstance(value, FunctionType): return BrinedFunction(value) elif isinstance(value, (list, tuple)): # create a duplicate of the collection with brined internals ty = type(value) return ty(brine(i) for i in iter(value)) elif isinstance(value, dict): items = value.items() return dict((key,brine(val)) for key,val in items) else: return value
[docs]def unbrine(value, with_globals=None): """ Unwrap a `value` previously wrapped with the `brine` function. Behavior by type of `value` is as follows: * `BrinedPartial` unwraps to a `partial` * `BrinedMethod` unwraps to a `instancemethod` * `BrinedFunction` unwraps to a `function` * `list` and `tuple` are duplicated and contents unbrined * `dict` is duplicated and its values are unbrined Parameters ---------- value : `object` value to be unbrined with_globals : `dict` or `None` globals dictionary to use when recreating functions. `None` is the same as `globals()` Returns ------- unwrapped : `object` as defined by the type of `value` """ glbls = globals() if with_globals is None else with_globals if isinstance(value, BrinedObject): return value.get(glbls) elif isinstance(value, (list, tuple)): ty = type(value) return ty(unbrine(i, glbls) for i in iter(value)) elif isinstance(value, dict): items = value.items() return dict((key,unbrine(val, glbls)) for key,val in items) else: return value
def code_unnew(code): """ The necessary arguments for use in :func:`code_new` to create an identical but distinct :data:`~types.CodeType` instance. :param code: code object for inspection :type code: :data:`~types.CodeType` :return: :class:`list` of member values of `code` """ return [ code.co_argcount, code.co_nlocals, code.co_stacksize, code.co_flags, code.co_code, code.co_consts, code.co_names, code.co_varnames, code.co_filename, code.co_name, code.co_firstlineno, code.co_lnotab, code.co_freevars, code.co_cellvars ] def code_new(argcount, nlocals, stacksize, flags, code, consts, names, varnames, filename, name, firstlineno, lnotab, freevars, cellvars): """ Create a new code object. Identical to :func:`new.code` :return: new :data:`~types.CodeType` instance """ return new.code(argcount, nlocals, stacksize, flags, code, consts, names, varnames, filename, name, firstlineno, lnotab, freevars, cellvars) def function_unnew(func): """ The necessary arguments for use in :func:`function_new` to create an identical but distinct :data:`~types.FunctionType` instance. :param func: function object for inspection :type func: :data:`~types.FunctionType` :return: :class:`list` of member values of `func` """ return [ func.func_code, func.func_globals, func.func_name, func.func_defaults, func.func_closure ] def function_new(code, with_globals, name, defaults, closure): """ Creates a new function. Identical to :func:`new.function` :param code: code object for the function to execute when called :type code: :data:`~types.CodeType` :param with_globals: :type with_globals: :class:`dict` often the result of :func:`globals` :param name: name for the function :type name: :class:`str` :param defaults: defaults for the function :param closure: captured :class:`CellType` cells defining the closure :type closure: :class:`tuple` :return: new :data:`~types.FunctionType` instance """ return new.function(code, with_globals, name, defaults, closure)
[docs]class BrinedObject(object): # pragma: no cover """ Abstract base class for brine wrappers. Defines the interface required for brine. * When instantiated, should accept a `value` to wrap * The `get` method should return a new copy of the wrapped `value` * Must define the `__getstate__` and `__setstate__` methods for `pickle` support. """ __metaclass__ = ABCMeta @abstractmethod
[docs] def __init__(self, value): """ Wrap `value` so that it may be pickled. Parameters ---------- value : `object` the instance needing to be wrapped for pickling """ pass
@abstractmethod
[docs] def __getstate__(self): """ Special method used by a `pickle.Pickler` to serialize an instance of this class. Returns ------- data : `tuple` information to provide to the pickle API. Should be the same as what is expected in the `__setstate__` method """ pass
@abstractmethod
[docs] def __setstate__(self, data): """ Special method used by a `pickle.Unpickler` to deserialize an instance of this class. Should accept the same data returned by the `__getstate__` method Parameters ---------- data : `tuple` information provided by the pickle API to instantiate a new copy """ pass
@abstractmethod
[docs] def get(self, with_globals): """ Recreate a copy of the initial brined value Parameters ---------- with_globals : `globals()` or `dict`-like The global namespace to be used when recreating the value. Returns ------- value : `object` A copy of the `value` originally passed to `__init__` """ pass # A function object needs to be brined before it can be pickled, and # unbrined after it's unpickled. We need to do this because pickle has # some default behavior for pickling types.FunctionType which we do # not want to break. Therefore, we will simply wrap any Function # instances in BrinedFunction before pickling, and unwap them after # unpickling
[docs]class BrinedFunction(BrinedObject): """ Wraps a function so that it may be pickled. For the most part you'll want to use the brine and unbrine functions from this module rather than instantiating or accessing this class directly """ def __init__(self, function): self._unfunc = self._function_unnew(function) self._fdict = dict(function.__dict__) def __getstate__(self): return self._unfunc, self._fdict def __setstate__(self, state): self._unfunc, self._fdict = state def _function_unnew(self, function): unfunc = function_unnew(function) unfunc[0] = self._code_unnew(unfunc[0]) unfunc[1] = dict() # strip out func_globals return unfunc def _code_unnew(self, code): """ can be overridden to process the unnew data """ return code_unnew(code) def get(self, with_globals): # compose the function func = self._function_new(with_globals, list(self._unfunc)) # setup any of the function's members func.__dict__.update(self._fdict) return func def _function_new(self, with_globals, ufunc): ufunc[0] = self._code_new(with_globals, list(ufunc[0])) ufunc[1] = with_globals return function_new(*ufunc) def _code_new(self, with_globals, uncode): return code_new(*uncode)
[docs]class BrinedMethod(BrinedObject): """ Wraps a bound method so that it can be pickled. By default pickle refuses to operate on bound instance method object. This wrapper will still require that the object instance supports pickling, which in turn requires that the class be defined at the top level. As with the `BrineFunction` class, it is better to use the `brine` and `unbrine` functions from this module rather than to create an instance of this class directly. """ def __init__(self, boundmethod): self._im_self = boundmethod.im_self self._funcname = boundmethod.im_func.__name__ def __getstate__(self): return (self._im_self, self._funcname) def __setstate__(self, state): self._im_self = state[0] self._funcname = state[1] def get(self, with_globals): return getattr(self._im_self, self._funcname)
[docs]class BrinedPartial(BrinedObject): """ Wrap a :class:`functools.partial` instance that references a function or method that is otherwise unsupported by pickle. """ def __init__(self, part): self._func = brine(part.func) self._args = brine(part.args or None) self._keywords = brine(part.keywords or None) def get(self, with_globals): func = unbrine(self._func, with_globals) args = unbrine(self._args or tuple(), with_globals) kwds = unbrine(self._keywords or dict(), with_globals) return partial(func, *args, **kwds) def __getstate__(self): return (self._func, self._args, self._keywords) def __setstate__(self, data): self._func, self._args, self._keywords = data # let's give the pickle module knowledge of how to load and dump Cell # and Code objects
def _pickle_cell(cell): return _unpickle_cell, (cell_get_value(cell), ) def _unpickle_cell(cell_val): return cell_from_value(cell_val) def reg_cell_pickler(): """ Called automatically when the module is loaded, this function will ensure that the CellType has pickle/unpickle functions registered with copy_reg """ copy_reg.pickle(CellType, _pickle_cell, _unpickle_cell) # register when the module is loaded reg_cell_pickler() def _pickle_code(code): return _unpickle_code, tuple(code_unnew(code)) def _unpickle_code(*code_val): return code_new(*code_val) def reg_code_pickler(): """ Called automatically when the module is loaded, this function will ensure that the CodeType has pickle/unpickle functions registered with copy_reg """ copy_reg.pickle(CodeType, _pickle_code, _unpickle_code) # Register when the module is loaded. Note, we do this to support # pickling of functions which return other functions. The code for the # anonymous inner is captured in a closure cell of the outer. reg_code_pickler() # # The end.