From 161841f46fbc28b88a0ad546c8dd482e5af1cb3e 2014-05-01 00:34:56
From: Thomas Kluyver <takowl@gmail.com>
Date: 2014-05-01 00:34:56
Subject: [PATCH] Merge pull request #5762 from minrk/better_closure_check

Fix check for pickling closures
---

diff --git a/IPython/utils/codeutil.py b/IPython/utils/codeutil.py
index a51af71..b3ffd6a 100644
--- a/IPython/utils/codeutil.py
+++ b/IPython/utils/codeutil.py
@@ -10,19 +10,6 @@ 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
-#-------------------------------------------------------------------------------
-
 import sys
 import types
 try:
@@ -34,12 +21,10 @@ 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)
diff --git a/IPython/utils/pickleutil.py b/IPython/utils/pickleutil.py
index d3fed87..bd69f40 100644
--- a/IPython/utils/pickleutil.py
+++ b/IPython/utils/pickleutil.py
@@ -2,18 +2,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
@@ -35,9 +25,11 @@ from IPython.config import Application
 if py3compat.PY3:
     buffer = memoryview
     class_type = type
+    closure_attr = '__closure__'
 else:
     from types import ClassType
     class_type = (type, ClassType)
+    closure_attr = 'func_closure'
 
 #-------------------------------------------------------------------------------
 # Functions
@@ -140,6 +132,9 @@ class CannedFunction(CannedObject):
             self.defaults = [ can(fd) for fd in f.__defaults__ ]
         else:
             self.defaults = None
+            
+        if getattr(f, closure_attr, None):
+            raise ValueError("Sorry, cannot pickle functions with closures.")
         self.module = f.__module__ or '__main__'
         self.__name__ = f.__name__
         self.buffers = []
diff --git a/IPython/utils/tests/test_pickleutil.py b/IPython/utils/tests/test_pickleutil.py
new file mode 100644
index 0000000..626e025
--- /dev/null
+++ b/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
+    
+    # true closures are not supported
+    with nt.assert_raises(ValueError):
+        pfoo = dumps(foo)
+
+    
\ No newline at end of file