IPython/utils/pickleutil.py

# encoding: utf-8

"""Pickle related utilities. Perhaps this should be called 'can'."""

__docformat__ = "restructuredtext en"

#-------------------------------------------------------------------------------
#  Copyright (C) 2008  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.
#-------------------------------------------------------------------------------

#-------------------------------------------------------------------------------
# Imports
#-------------------------------------------------------------------------------

import copy
import sys
from types import FunctionType

import codeutil

#-------------------------------------------------------------------------------
# Classes
#-------------------------------------------------------------------------------


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)))
            
    
    def getObject(self, g=None):
        if g is None:
            g = globals()
        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
    
    def __repr__(self):
        return "<Reference: %r>"%self.name
    
    def getObject(self, g=None):
        if g is None:
            g = globals()
        try:
            return g[self.name]
        except KeyError:
            raise NameError("name %r is not defined"%self.name)
    

class CannedFunction(CannedObject):
    
    def __init__(self, f):
        self._checkType(f)    
        self.code = f.func_code
        self.defaults = f.func_defaults
        self.module = f.__module__ or '__main__'
        self.__name__ = f.__name__
    
    def _checkType(self, obj):
        assert isinstance(obj, FunctionType), "Not a function type"
    
    def getObject(self, g=None):
        # try to load function back into its module:
        if not self.module.startswith('__'):
            try:
                __import__(self.module)
            except ImportError:
                pass
            else:
                g = sys.modules[self.module].__dict__
            
        if g is None:
            g = globals()
        newFunc = FunctionType(self.code, g, self.__name__, self.defaults)
        return newFunc

#-------------------------------------------------------------------------------
# Functions
#-------------------------------------------------------------------------------

def can(obj):
    # import here to prevent module-level circular imports
    from IPython.parallel import dependent
    if isinstance(obj, dependent):
        keys = ('f','df')
        return CannedObject(obj, keys=keys)
    elif isinstance(obj, FunctionType):
        return CannedFunction(obj)
    elif isinstance(obj,dict):
        return canDict(obj)
    elif isinstance(obj, (list,tuple)):
        return canSequence(obj)
    else:
        return obj

def canDict(obj):
    if isinstance(obj, dict):
        newobj = {}
        for k, v in obj.iteritems():
            newobj[k] = can(v)
        return newobj
    else:
        return obj

def canSequence(obj):
    if isinstance(obj, (list, tuple)):
        t = type(obj)
        return t([can(i) for i in obj])
    else:
        return obj

def uncan(obj, g=None):
    if isinstance(obj, CannedObject):
        return obj.getObject(g)
    elif isinstance(obj,dict):
        return uncanDict(obj, g)
    elif isinstance(obj, (list,tuple)):
        return uncanSequence(obj, g)
    else:
        return obj

def uncanDict(obj, g=None):
    if isinstance(obj, dict):
        newobj = {}
        for k, v in obj.iteritems():
            newobj[k] = uncan(v,g)
        return newobj
    else:
        return obj

def uncanSequence(obj, g=None):
    if isinstance(obj, (list, tuple)):
        t = type(obj)
        return t([uncan(i,g) for i in obj])
    else:
        return obj


def rebindFunctionGlobals(f, glbls):
    return FunctionType(f.func_code, glbls)