IPython/zmq/newserialized.py

# encoding: utf-8
# -*- test-case-name: IPython.kernel.test.test_newserialized -*-

"""Refactored serialization classes and interfaces."""

__docformat__ = "restructuredtext en"

# Tell nose to skip this module
__test__ = {}

#-------------------------------------------------------------------------------
#  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 cPickle as pickle

# from twisted.python import components
# from zope.interface import Interface, implements

try:
    import numpy
except ImportError:
    pass

from IPython.kernel.error import SerializationError

#-----------------------------------------------------------------------------
# Classes and functions
#-----------------------------------------------------------------------------

class ISerialized:
    
    def getData():
        """"""

    def getDataSize(units=10.0**6):
        """"""

    def getTypeDescriptor():
        """"""

    def getMetadata():
        """"""
        
        
class IUnSerialized:
    
    def getObject():
        """"""
        
class Serialized(object):
    
    # implements(ISerialized)
    
    def __init__(self, data, typeDescriptor, metadata={}):
        self.data = data
        self.typeDescriptor = typeDescriptor
        self.metadata = metadata
        
    def getData(self):
        return self.data
        
    def getDataSize(self, units=10.0**6):
        return len(self.data)/units
        
    def getTypeDescriptor(self):
        return self.typeDescriptor
        
    def getMetadata(self):
        return self.metadata

        
class UnSerialized(object):
    
    # implements(IUnSerialized)
        
    def __init__(self, obj):
        self.obj = obj
        
    def getObject(self):
        return self.obj

        
class SerializeIt(object):
    
    # implements(ISerialized)
    
    def __init__(self, unSerialized):
        self.data = None
        self.obj = unSerialized.getObject()
        if globals().has_key('numpy') and isinstance(self.obj, numpy.ndarray):
                if len(self.obj) == 0:         # length 0 arrays can't be reconstructed
                    raise SerializationError("You cannot send a length 0 array")
                self.obj = numpy.ascontiguousarray(self.obj, dtype=None)
                self.typeDescriptor = 'ndarray'
                self.metadata = {'shape':self.obj.shape,
                                 'dtype':self.obj.dtype.str}
        elif isinstance(self.obj, str):
            self.typeDescriptor = 'bytes'
            self.metadata = {}
        elif isinstance(self.obj, buffer):
            self.typeDescriptor = 'buffer'
            self.metadata = {}
        else:
            self.typeDescriptor = 'pickle'
            self.metadata = {}
        self._generateData()            
    
    def _generateData(self):
        if self.typeDescriptor == 'ndarray':
            self.data = numpy.getbuffer(self.obj)
        elif self.typeDescriptor in ('bytes', 'buffer'):
            self.data = self.obj
        elif self.typeDescriptor == 'pickle':
            self.data = pickle.dumps(self.obj, -1)
        else:
            raise SerializationError("Really wierd serialization error.")
        del self.obj
        
    def getData(self):
        return self.data
        
    def getDataSize(self, units=10.0**6):
        return 1.0*len(self.data)/units
        
    def getTypeDescriptor(self):
        return self.typeDescriptor
        
    def getMetadata(self):
        return self.metadata


class UnSerializeIt(UnSerialized):
    
    # implements(IUnSerialized)
    
    def __init__(self, serialized):
        self.serialized = serialized
        
    def getObject(self):
        typeDescriptor = self.serialized.getTypeDescriptor()
        if globals().has_key('numpy') and typeDescriptor == 'ndarray':
                result = numpy.frombuffer(self.serialized.getData(), dtype = self.serialized.metadata['dtype'])
                result.shape = self.serialized.metadata['shape']
                # This is a hack to make the array writable.  We are working with
                # the numpy folks to address this issue.
                result = result.copy()
        elif typeDescriptor == 'pickle':
            result = pickle.loads(self.serialized.getData())
        elif typeDescriptor in ('bytes', 'buffer'):
            result = self.serialized.getData()
        else:
            raise SerializationError("Really wierd serialization error.")
        return result

def serialize(obj):
    return SerializeIt(UnSerialized(obj))
    
def unserialize(serialized):
    return UnSerializeIt(serialized).getObject()