adjust how canning deals with import strings

* imports should only happen once
* errors should only be logged when *user* imports fail (not no-numpy)

IPython/utils/pickleutil.py

@@ -1,244 +1,281 b''
 # 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.
 #-------------------------------------------------------------------------------
 
 #-------------------------------------------------------------------------------
 # Imports
 #-------------------------------------------------------------------------------
 
 import copy
 import logging
 import sys
 from types import FunctionType
 
 try:
     import cPickle as pickle
 except ImportError:
     import pickle
 
 try:
     import numpy
 except:
     numpy = None
 
 import codeutil
 import py3compat
 from importstring import import_item
 
 from IPython.config import Application
 
 if py3compat.PY3:
     buffer = memoryview
 
 #-------------------------------------------------------------------------------
 # 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)))
         
         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('__'):
             try:
                 __import__(self.module)
             except ImportError:
                 pass
             else:
                 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 _error(*args, **kwargs):
+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
     else:
         logger = logging.getLogger()
         if not logger.handlers:
             logging.basicConfig()
-    logger.error(*args, **kwargs)
+    
+    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):
-            try:
-                cls = import_item(cls)
-            except Exception:
-                _error("cannning class not importable: %r", cls, exc_info=True)
-                cls = None
-                continue
-        if isinstance(obj, cls):
+            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
     else:
         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])
     else:
         return obj
 
 def uncan(obj, g=None):
     """invert canning"""
+    
+    import_needed = False
     for cls,uncanner in uncan_map.iteritems():
         if isinstance(cls, basestring):
-            try:
-                cls = import_item(cls)
-            except Exception:
-                _error("uncanning class not importable: %r", cls, exc_info=True)
-                cls = None
-                continue
-        if isinstance(obj, cls):
+            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
     else:
         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])
     else:
         return obj
 
 
 #-------------------------------------------------------------------------------
-# API dictionary
+# 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()