Merge pull request #5761 from minrk/pickle_simple_closures

handle simple closures in pickleutil

IPython/utils/tests/test_pickleutil.py

@@ -0,0 +1,62
+
+import pickle
+
+import nose.tools as nt
+from IPython.utils import codeutil
+from IPython.utils.pickleutil import can, uncan
+
+def interactive(f):
+    f.__module__ = '__main__'
+    return f
+
+def dumps(obj):
+    return pickle.dumps(can(obj))
+
+def loads(obj):
+    return uncan(pickle.loads(obj))
+
+def test_no_closure():
+    @interactive
+    def foo():
+        a = 5
+        return a
+    
+    pfoo = dumps(foo)
+    bar = loads(pfoo)
+    nt.assert_equal(foo(), bar())
+
+def test_generator_closure():
+    # this only creates a closure on Python 3
+    @interactive
+    def foo():
+        i = 'i'
+        r = [ i for j in (1,2) ]
+        return r
+    
+    pfoo = dumps(foo)
+    bar = loads(pfoo)
+    nt.assert_equal(foo(), bar())
+
+def test_nested_closure():
+    @interactive
+    def foo():
+        i = 'i'
+        def g():
+            return i
+        return g()
+    
+    pfoo = dumps(foo)
+    bar = loads(pfoo)
+    nt.assert_equal(foo(), bar())
+
+def test_closure():
+    i = 'i'
+    @interactive
+    def foo():
+        return i
+    
+    pfoo = dumps(foo)
+    bar = loads(pfoo)
+    nt.assert_equal(foo(), bar())
+
+    
  No newline at end of file

IPython/utils/codeutil.py

@@ -1,47 +1,35
 # encoding: utf-8
 
 """Utilities to enable code objects to be pickled.
 
 Any process that import this module will be able to pickle code objects.  This
 includes the func_code attribute of any function.  Once unpickled, new
 functions can be built using new.function(code, globals()).  Eventually
 we need to automate all of this so that functions themselves can be pickled.
 
 Reference: A. Tremols, P Cogolo, "Python Cookbook," p 302-305
 """
 
-__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
-#-------------------------------------------------------------------------------
+# Copyright (c) IPython Development Team.
+# Distributed under the terms of the Modified BSD License.
 
 import sys
 import types
 try:
     import copyreg  # Py 3
 except ImportError:
     import copy_reg as copyreg  # Py 2
 
 def code_ctor(*args):
     return types.CodeType(*args)
 
 def reduce_code(co):
-    if co.co_freevars or co.co_cellvars:
-        raise ValueError("Sorry, cannot pickle code objects with closures")
     args =  [co.co_argcount, co.co_nlocals, co.co_stacksize,
             co.co_flags, co.co_code, co.co_consts, co.co_names,
             co.co_varnames, co.co_filename, co.co_name, co.co_firstlineno,
-            co.co_lnotab]
+            co.co_lnotab, co.co_freevars, co.co_cellvars]
     if sys.version_info[0] >= 3:
         args.insert(1, co.co_kwonlyargcount)
     return code_ctor, tuple(args)
 
 copyreg.pickle(types.CodeType, reduce_code)
  No newline at end of file

IPython/utils/pickleutil.py

@@ -1,390 +1,413
 # 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
-#-------------------------------------------------------------------------------
+# Copyright (c) IPython Development Team.
+# Distributed under the terms of the Modified BSD License.
 
 import copy
 import logging
 import sys
 from types import FunctionType
 
 try:
     import cPickle as pickle
 except ImportError:
     import pickle
 
 from . import codeutil  # This registers a hook when it's imported
 from . import py3compat
 from .importstring import import_item
 from .py3compat import string_types, iteritems
 
 from IPython.config import Application
 
 if py3compat.PY3:
     buffer = memoryview
     class_type = type
 else:
     from types import ClassType
     class_type = (type, ClassType)
 
+def _get_cell_type(a=None):
+    """the type of a closure cell doesn't seem to be importable,
+    so just create one
+    """
+    def inner():
+        return a
+    return type(py3compat.get_closure(inner)[0])
+
+cell_type = _get_cell_type()
+
 #-------------------------------------------------------------------------------
 # Functions
 #-------------------------------------------------------------------------------
 
 
 def use_dill():
     """use dill to expand serialization support
     
     adds support for object methods and closures to serialization.
     """
     # import dill causes most of the magic
     import dill
     
     # dill doesn't work with cPickle,
     # tell the two relevant modules to use plain pickle
     
     global pickle
     pickle = dill
 
     try:
         from IPython.kernel.zmq import serialize
     except ImportError:
         pass
     else:
         serialize.pickle = dill
     
     # disable special function handling, let dill take care of it
     can_map.pop(FunctionType, None)
 
 
 #-------------------------------------------------------------------------------
 # Classes
 #-------------------------------------------------------------------------------
 
 
 class CannedObject(object):
     def __init__(self, obj, keys=[], hook=None):
         """can an object for safe pickling
         
         Parameters
         ==========
         
         obj:
             The object to be canned
         keys: list (optional)
             list of attribute names that will be explicitly canned / uncanned
         hook: callable (optional)
             An optional extra callable,
             which can do additional processing of the uncanned object.
         
         large data may be offloaded into the buffers list,
         used for zero-copy transfers.
         """
         self.keys = keys
         self.obj = copy.copy(obj)
         self.hook = can(hook)
         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 = {}
         obj = self.obj
         for key in self.keys:
             setattr(obj, key, uncan(getattr(obj, key), g))
         
         if self.hook:
             self.hook = uncan(self.hook, g)
             self.hook(obj, g)
         return self.obj
     
 
 class Reference(CannedObject):
     """object for wrapping a remote reference by name."""
     def __init__(self, name):
         if not isinstance(name, string_types):
             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 CannedCell(CannedObject):
+    """Can a closure cell"""
+    def __init__(self, cell):
+        self.cell_contents = can(cell.cell_contents)
+    
+    def get_object(self, g=None):
+        cell_contents = uncan(self.cell_contents, g)
+        def inner():
+            return cell_contents
+        return py3compat.get_closure(inner)[0]
+
+
 class CannedFunction(CannedObject):
 
     def __init__(self, f):
         self._check_type(f)
         self.code = f.__code__
         if f.__defaults__:
             self.defaults = [ can(fd) for fd in f.__defaults__ ]
         else:
             self.defaults = None
+        
+        closure = py3compat.get_closure(f)
+        if closure:
+            self.closure = tuple( can(cell) for cell in closure )
+        else:
+            self.closure = 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)
+        if self.closure:
+            closure = tuple(uncan(cell, g) for cell in self.closure)
+        else:
+            closure = None
+        newFunc = FunctionType(self.code, g, self.__name__, defaults, closure)
         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, class_type), "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):
         from numpy import ascontiguousarray
         self.shape = obj.shape
         self.dtype = obj.dtype.descr if obj.dtype.fields else obj.dtype.str
         self.pickled = False
         if sum(obj.shape) == 0:
             self.pickled = True
         elif obj.dtype == 'O':
             # can't handle object dtype with buffer approach
             self.pickled = True
         elif obj.dtype.fields and any(dt == 'O' for dt,sz in obj.dtype.fields.values()):
             self.pickled = True
         if self.pickled:
             # just pickle it
             self.buffers = [pickle.dumps(obj, -1)]
         else:
             # ensure contiguous
             obj = ascontiguousarray(obj, dtype=None)
             self.buffers = [buffer(obj)]
     
     def get_object(self, g=None):
         from numpy import frombuffer
         data = self.buffers[0]
         if self.pickled:
             # no shape, we just pickled it
             return pickle.loads(data)
         else:
             return 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
     else:
         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 list(mapping.items()):
         if isinstance(key, string_types):
             try:
                 cls = import_item(key)
             except Exception:
                 if original and key not in original:
                     # only message on user-added classes
                     log.error("canning class not importable: %r", key, exc_info=True)
                 mapping.pop(key)
             else:
                 mapping[cls] = mapping.pop(key)
 
 def istype(obj, check):
     """like isinstance(obj, check), but strict
     
     This won't catch subclasses.
     """
     if isinstance(check, tuple):
         for cls in check:
             if type(obj) is cls:
                 return True
         return False
     else:
         return type(obj) is check
 
 def can(obj):
     """prepare an object for pickling"""
     
     import_needed = False
     
     for cls,canner in iteritems(can_map):
         if isinstance(cls, string_types):
             import_needed = True
             break
         elif istype(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_class(obj):
     if isinstance(obj, class_type) and obj.__module__ == '__main__':
         return CannedClass(obj)
     else:
         return obj
 
 def can_dict(obj):
     """can the *values* of a dict"""
     if istype(obj, dict):
         newobj = {}
         for k, v in iteritems(obj):
             newobj[k] = can(v)
         return newobj
     else:
         return obj
 
 sequence_types = (list, tuple, set)
 
 def can_sequence(obj):
     """can the elements of a sequence"""
     if istype(obj, sequence_types):
         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 iteritems(uncan_map):
         if isinstance(cls, string_types):
             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 istype(obj, dict):
         newobj = {}
         for k, v in iteritems(obj):
             newobj[k] = uncan(v,g)
         return newobj
     else:
         return obj
 
 def uncan_sequence(obj, g=None):
     if istype(obj, sequence_types):
         t = type(obj)
         return t([uncan(i,g) for i in obj])
     else:
         return obj
 
 def _uncan_dependent_hook(dep, g=None):
     dep.check_dependency()
     
 def can_dependent(obj):
     return CannedObject(obj, keys=('f', 'df'), hook=_uncan_dependent_hook)
 
 #-------------------------------------------------------------------------------
 # API dictionaries
 #-------------------------------------------------------------------------------
 
 # These dicts can be extended for custom serialization of new objects
 
 can_map = {
     'IPython.parallel.dependent' : can_dependent,
     'numpy.ndarray' : CannedArray,
     FunctionType : CannedFunction,
     bytes : CannedBytes,
     buffer : CannedBuffer,
+    cell_type : CannedCell,
     class_type : can_class,
 }
 
 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()

IPython/utils/py3compat.py

@@ -1,264 +1,271
 # coding: utf-8
 """Compatibility tricks for Python 3. Mainly to do with unicode."""
 import functools
 import os
 import sys
 import re
 import types
 
 from .encoding import DEFAULT_ENCODING
 
 orig_open = open
 
 def no_code(x, encoding=None):
     return x
 
 def decode(s, encoding=None):
     encoding = encoding or DEFAULT_ENCODING
     return s.decode(encoding, "replace")
 
 def encode(u, encoding=None):
     encoding = encoding or DEFAULT_ENCODING
     return u.encode(encoding, "replace")
 
 
 def cast_unicode(s, encoding=None):
     if isinstance(s, bytes):
         return decode(s, encoding)
     return s
 
 def cast_bytes(s, encoding=None):
     if not isinstance(s, bytes):
         return encode(s, encoding)
     return s
 
 def _modify_str_or_docstring(str_change_func):
     @functools.wraps(str_change_func)
     def wrapper(func_or_str):
         if isinstance(func_or_str, string_types):
             func = None
             doc = func_or_str
         else:
             func = func_or_str
             doc = func.__doc__
         
         doc = str_change_func(doc)
         
         if func:
             func.__doc__ = doc
             return func
         return doc
     return wrapper
 
 def safe_unicode(e):
     """unicode(e) with various fallbacks. Used for exceptions, which may not be
     safe to call unicode() on.
     """
     try:
         return unicode_type(e)
     except UnicodeError:
         pass
 
     try:
         return str_to_unicode(str(e))
     except UnicodeError:
         pass
 
     try:
         return str_to_unicode(repr(e))
     except UnicodeError:
         pass
 
     return u'Unrecoverably corrupt evalue'
 
 if sys.version_info[0] >= 3:
     PY3 = True
     
     # keep reference to builtin_mod because the kernel overrides that value
     # to forward requests to a frontend.
     def input(prompt=''):
         return builtin_mod.input(prompt)
     
     builtin_mod_name = "builtins"
     import builtins as builtin_mod
     
     str_to_unicode = no_code
     unicode_to_str = no_code
     str_to_bytes = encode
     bytes_to_str = decode
     cast_bytes_py2 = no_code
     cast_unicode_py2 = no_code
     
     string_types = (str,)
     unicode_type = str
     
     def isidentifier(s, dotted=False):
         if dotted:
             return all(isidentifier(a) for a in s.split("."))
         return s.isidentifier()
     
     open = orig_open
     xrange = range
     def iteritems(d): return iter(d.items())
     def itervalues(d): return iter(d.values())
     getcwd = os.getcwd
     
     MethodType = types.MethodType
     
     def execfile(fname, glob, loc=None):
         loc = loc if (loc is not None) else glob
         with open(fname, 'rb') as f:
             exec(compile(f.read(), fname, 'exec'), glob, loc)
     
     # Refactor print statements in doctests.
     _print_statement_re = re.compile(r"\bprint (?P<expr>.*)$", re.MULTILINE)
     def _print_statement_sub(match):
         expr = match.groups('expr')
         return "print(%s)" % expr
     
     @_modify_str_or_docstring
     def doctest_refactor_print(doc):
         """Refactor 'print x' statements in a doctest to print(x) style. 2to3
         unfortunately doesn't pick up on our doctests.
         
         Can accept a string or a function, so it can be used as a decorator."""
         return _print_statement_re.sub(_print_statement_sub, doc)
     
     # Abstract u'abc' syntax:
     @_modify_str_or_docstring
     def u_format(s):
         """"{u}'abc'" --> "'abc'" (Python 3)
         
         Accepts a string or a function, so it can be used as a decorator."""
         return s.format(u='')
+    
+    def get_closure(f):
+        """Get a function's closure attribute"""
+        return f.__closure__
 
 else:
     PY3 = False
     
     # keep reference to builtin_mod because the kernel overrides that value
     # to forward requests to a frontend.
     def input(prompt=''):
         return builtin_mod.raw_input(prompt)
     
     builtin_mod_name = "__builtin__"
     import __builtin__ as builtin_mod
     
     str_to_unicode = decode
     unicode_to_str = encode
     str_to_bytes = no_code
     bytes_to_str = no_code
     cast_bytes_py2 = cast_bytes
     cast_unicode_py2 = cast_unicode
     
     string_types = (str, unicode)
     unicode_type = unicode
     
     import re
     _name_re = re.compile(r"[a-zA-Z_][a-zA-Z0-9_]*$")
     def isidentifier(s, dotted=False):
         if dotted:
             return all(isidentifier(a) for a in s.split("."))
         return bool(_name_re.match(s))
     
     class open(object):
         """Wrapper providing key part of Python 3 open() interface."""
         def __init__(self, fname, mode="r", encoding="utf-8"):
             self.f = orig_open(fname, mode)
             self.enc = encoding
         
         def write(self, s):
             return self.f.write(s.encode(self.enc))
         
         def read(self, size=-1):
             return self.f.read(size).decode(self.enc)
         
         def close(self):
             return self.f.close()
         
         def __enter__(self):
             return self
         
         def __exit__(self, etype, value, traceback):
             self.f.close()
     
     xrange = xrange
     def iteritems(d): return d.iteritems()
     def itervalues(d): return d.itervalues()
     getcwd = os.getcwdu
 
     def MethodType(func, instance):
         return types.MethodType(func, instance, type(instance))
     
     def doctest_refactor_print(func_or_str):
         return func_or_str
 
+    def get_closure(f):
+        """Get a function's closure attribute"""
+        return f.func_closure
 
     # Abstract u'abc' syntax:
     @_modify_str_or_docstring
     def u_format(s):
         """"{u}'abc'" --> "u'abc'" (Python 2)
         
         Accepts a string or a function, so it can be used as a decorator."""
         return s.format(u='u')
 
     if sys.platform == 'win32':
         def execfile(fname, glob=None, loc=None):
             loc = loc if (loc is not None) else glob
             # The rstrip() is necessary b/c trailing whitespace in files will
             # cause an IndentationError in Python 2.6 (this was fixed in 2.7,
             # but we still support 2.6).  See issue 1027.
             scripttext = builtin_mod.open(fname).read().rstrip() + '\n'
             # compile converts unicode filename to str assuming
             # ascii. Let's do the conversion before calling compile
             if isinstance(fname, unicode):
                 filename = unicode_to_str(fname)
             else:
                 filename = fname
             exec(compile(scripttext, filename, 'exec'), glob, loc)
     else:
         def execfile(fname, *where):
             if isinstance(fname, unicode):
                 filename = fname.encode(sys.getfilesystemencoding())
             else:
                 filename = fname
             builtin_mod.execfile(filename, *where)
 
 
 def annotate(**kwargs):
     """Python 3 compatible function annotation for Python 2."""
     if not kwargs:
         raise ValueError('annotations must be provided as keyword arguments')
     def dec(f):
         if hasattr(f, '__annotations__'):
             for k, v in kwargs.items():
                 f.__annotations__[k] = v
         else:
             f.__annotations__ = kwargs
         return f
     return dec
 
 
 # Parts below taken from six:
 # Copyright (c) 2010-2013 Benjamin Peterson
 #
 # Permission is hereby granted, free of charge, to any person obtaining a copy
 # of this software and associated documentation files (the "Software"), to deal
 # in the Software without restriction, including without limitation the rights
 # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 # copies of the Software, and to permit persons to whom the Software is
 # furnished to do so, subject to the following conditions:
 #
 # The above copyright notice and this permission notice shall be included in all
 # copies or substantial portions of the Software.
 #
 # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
 def with_metaclass(meta, *bases):
     """Create a base class with a metaclass."""
     return meta("_NewBase", bases, {})