iframe-proxy | Sunbelt Computer Software

277 lines (219 loc) · 7.46 KB
# encoding: utf-8
"""Pickle related utilities. Perhaps this should be called 'can'."""
__docformat__ = "restructuredtext en"
#-------------------------------------------------------------------------------
#  Copyright (C) 2008-2011  The IPython Development Team
#  Distributed under the terms of the BSD License.  The full license is in
#  the file COPYING, distributed as part of this software.
#-------------------------------------------------------------------------------
#-------------------------------------------------------------------------------
#-------------------------------------------------------------------------------
import copy
import logging
from types import FunctionType
    import cPickle as pickle
except ImportError:
    import pickle
    import numpy
    numpy = None
import codeutil
import py3compat
from importstring import import_item
from IPython.config import Application
if py3compat.PY3:
    buffer = memoryview
#-------------------------------------------------------------------------------
#-------------------------------------------------------------------------------
class CannedObject(object):
    def __init__(self, obj, keys=[]):
        self.keys = keys
        self.obj = copy.copy(obj)
        for key in keys:
            setattr(self.obj, key, can(getattr(obj, key)))
        self.buffers = []
    def get_object(self, g=None):
        if g is None:
            g = {}
        for key in self.keys:
            setattr(self.obj, key, uncan(getattr(self.obj, key), g))
        return self.obj
class Reference(CannedObject):
    """object for wrapping a remote reference by name."""
    def __init__(self, name):
        if not isinstance(name, basestring):
            raise TypeError("illegal name: %r"%name)
        self.name = name
        self.buffers = []
    def __repr__(self):
        return "<Reference: %r>"%self.name
    def get_object(self, g=None):
        if g is None:
            g = {}
        return eval(self.name, g)
class CannedFunction(CannedObject):
    def __init__(self, f):
        self._check_type(f)
        self.code = f.func_code
        if f.func_defaults:
            self.defaults = [ can(fd) for fd in f.func_defaults ]
        else:
            self.defaults = None
        self.module = f.__module__ or '__main__'
        self.__name__ = f.__name__
        self.buffers = []
    def _check_type(self, obj):
        assert isinstance(obj, FunctionType), "Not a function type"
    def get_object(self, g=None):
        # try to load function back into its module:
        if not self.module.startswith('__'):
            __import__(self.module)
            g = sys.modules[self.module].__dict__
        if g is None:
            g = {}
        if self.defaults:
            defaults = tuple(uncan(cfd, g) for cfd in self.defaults)
        else:
            defaults = None
        newFunc = FunctionType(self.code, g, self.__name__, defaults)
        return newFunc
class CannedArray(CannedObject):
    def __init__(self, obj):
        self.shape = obj.shape
        self.dtype = obj.dtype.descr if obj.dtype.fields else obj.dtype.str
        if sum(obj.shape) == 0:
            # just pickle it
            self.buffers = [pickle.dumps(obj, -1)]
        else:
            # ensure contiguous
            obj = numpy.ascontiguousarray(obj, dtype=None)
            self.buffers = [buffer(obj)]
    def get_object(self, g=None):
        data = self.buffers[0]
        if sum(self.shape) == 0:
            # no shape, we just pickled it
            return pickle.loads(data)
        else:
            return numpy.frombuffer(data, dtype=self.dtype).reshape(self.shape)
class CannedBytes(CannedObject):
    wrap = bytes
    def __init__(self, obj):
        self.buffers = [obj]
    def get_object(self, g=None):
        data = self.buffers[0]
        return self.wrap(data)
def CannedBuffer(CannedBytes):
    wrap = buffer
#-------------------------------------------------------------------------------
# Functions
#-------------------------------------------------------------------------------
def _logger():
    """get the logger for the current Application
    the root logger will be used if no Application is running
    if Application.initialized():
        logger = Application.instance().log
        logger = logging.getLogger()
        if not logger.handlers:
            logging.basicConfig()
    return logger
def _import_mapping(mapping, original=None):
    """import any string-keys in a type mapping
    log = _logger()
    log.debug("Importing canning map")
    for key,value in mapping.items():
        if isinstance(key, basestring):
            try:
                cls = import_item(key)
            except Exception:
                if original and key not in original:
                    # only message on user-added classes
                    log.error("cannning class not importable: %r", key, exc_info=True)
                mapping.pop(key)
            else:
                mapping[cls] = mapping.pop(key)
def can(obj):
    """prepare an object for pickling"""
    import_needed = False
    for cls,canner in can_map.iteritems():
        if isinstance(cls, basestring):
            import_needed = True
            break
        elif isinstance(obj, cls):
            return canner(obj)
    if import_needed:
        # perform can_map imports, then try again
        # this will usually only happen once
        _import_mapping(can_map, _original_can_map)
        return can(obj)
    return obj
def can_dict(obj):
    """can the *values* of a dict"""
    if isinstance(obj, dict):
        newobj = {}
        for k, v in obj.iteritems():
            newobj[k] = can(v)
        return newobj
        return obj
def can_sequence(obj):
    """can the elements of a sequence"""
    if isinstance(obj, (list, tuple)):
        t = type(obj)
        return t([can(i) for i in obj])
        return obj
def uncan(obj, g=None):
    """invert canning"""
    import_needed = False
    for cls,uncanner in uncan_map.iteritems():
        if isinstance(cls, basestring):
            import_needed = True
            break
        elif isinstance(obj, cls):
            return uncanner(obj, g)
    if import_needed:
        # perform uncan_map imports, then try again
        # this will usually only happen once
        _import_mapping(uncan_map, _original_uncan_map)
        return uncan(obj, g)
    return obj
def uncan_dict(obj, g=None):
    if isinstance(obj, dict):
        newobj = {}
        for k, v in obj.iteritems():
            newobj[k] = uncan(v,g)
        return newobj
        return obj
def uncan_sequence(obj, g=None):
    if isinstance(obj, (list, tuple)):
        t = type(obj)
        return t([uncan(i,g) for i in obj])
        return obj
#-------------------------------------------------------------------------------
# API dictionaries
#-------------------------------------------------------------------------------
# These dicts can be extended for custom serialization of new objects
can_map = {
    'IPython.parallel.dependent' : lambda obj: CannedObject(obj, keys=('f','df')),
    'numpy.ndarray' : CannedArray,
    FunctionType : CannedFunction,
    bytes : CannedBytes,
    buffer : CannedBuffer,
uncan_map = {
    CannedObject : lambda obj, g: obj.get_object(g),
# for use in _import_mapping:
_original_can_map = can_map.copy()
_original_uncan_map = uncan_map.copy()
Sunbelt Computer Software

PL/B Language Development and Support

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Sunbelt Computer Software

PL/B Language Development and Support

FilesExpand file tree

pickleutil.py

Latest commit

History

pickleutil.py

File metadata and controls
