From ba376d069fdd618aced799337997e3ea3d0c46f7 2013-01-07 01:04:16
From: MinRK <benjaminrk@gmail.com>
Date: 2013-01-07 01:04:16
Subject: [PATCH] can classes
---

diff --git a/IPython/utils/pickleutil.py b/IPython/utils/pickleutil.py
index 3b4735a..01e8583 100644
--- a/IPython/utils/pickleutil.py
+++ b/IPython/utils/pickleutil.py
@@ -18,7 +18,7 @@ __docformat__ = "restructuredtext en"
 import copy
 import logging
 import sys
-from types import FunctionType
+from types import FunctionType, ClassType
 
 try:
     import cPickle as pickle
@@ -110,6 +110,30 @@ class CannedFunction(CannedObject):
         newFunc = FunctionType(self.code, g, self.__name__, defaults)
         return newFunc
 
+class CannedClass(CannedObject):
+
+    def __init__(self, cls):
+        self._check_type(cls)
+        self.name = cls.__name__
+        self.old_style = not isinstance(cls, type)
+        self._canned_dict = {}
+        for k,v in cls.__dict__.items():
+            if k not in ('__weakref__', '__dict__'):
+                self._canned_dict[k] = can(v)
+        if self.old_style:
+            mro = []
+        else:
+            mro = cls.mro()
+        
+        self.parents = [ can(c) for c in mro[1:] ]
+        self.buffers = []
+
+    def _check_type(self, obj):
+        assert isinstance(obj, (type, ClassType)), "Not a class type"
+
+    def get_object(self, g=None):
+        parents = tuple(uncan(p, g) for p in self.parents)
+        return type(self.name, parents, uncan_dict(self._canned_dict, g=g))
 
 class CannedArray(CannedObject):
     def __init__(self, obj):
@@ -200,6 +224,12 @@ def can(obj):
     
     return obj
 
+def can_class(obj):
+    if isinstance(obj, (type, ClassType)) and obj.__module__ == '__main__':
+        return CannedClass(obj)
+    else:
+        return obj
+
 def can_dict(obj):
     """can the *values* of a dict"""
     if isinstance(obj, dict):
@@ -266,6 +296,8 @@ can_map = {
     FunctionType : CannedFunction,
     bytes : CannedBytes,
     buffer : CannedBuffer,
+    type : can_class,
+    ClassType : can_class,
 }
 
 uncan_map = {