From b4609459482154bf50410a75d6a31f7cd7165fbe 2011-11-27 00:57:32
From: Thomas Kluyver <takowl@gmail.com>
Date: 2011-11-27 00:57:32
Subject: [PATCH] Use user_ns as global namespace if it is passed without user_module, and add test for pickling interactively defined objects.

Closes gh-29

---

diff --git a/IPython/core/interactiveshell.py b/IPython/core/interactiveshell.py
index 5875e4b..adf522d 100644
--- a/IPython/core/interactiveshell.py
+++ b/IPython/core/interactiveshell.py
@@ -888,19 +888,13 @@ class InteractiveShell(SingletonConfigurable, Magic):
         # should start with "import __builtin__" (note, no 's') which will
         # definitely give you a module. Yeah, it's somewhat confusing:-(.
 
-        # These routines return properly built dicts as needed by the rest of
-        # the code, and can also be used by extension writers to generate
-        # properly initialized namespaces.
-        self.user_module = self.prepare_user_module(user_module)
+        # These routines return a properly built module and dict as needed by
+        # the rest of the code, and can also be used by extension writers to
+        # generate properly initialized namespaces.
+        self.user_module, self.user_ns = self.prepare_user_module(user_module, user_ns)
 
-        if user_ns is None:
-            user_ns = self.user_module.__dict__
-        self.user_ns = user_ns
-
-        # An auxiliary namespace that checks what parts of the user_ns were
-        # loaded at startup, so we can list later only variables defined in
-        # actual interactive use.  Since it is always a subset of user_ns, it
-        # doesn't need to be separately tracked in the ns_table.
+        # A record of hidden variables we have added to the user namespace, so
+        # we can list later only variables defined in actual interactive use.
         self.user_ns_hidden = set()
 
         # Now that FakeModule produces a real module, we've run into a nasty
@@ -943,20 +937,38 @@ class InteractiveShell(SingletonConfigurable, Magic):
     def user_global_ns(self):
         return self.user_module.__dict__
 
-    def prepare_user_module(self, user_module=None):
-        """Prepares a module for use as the interactive __main__ module in
-        which user code is run.
+    def prepare_user_module(self, user_module=None, user_ns=None):
+        """Prepare the module and namespace in which user code will be run.
+        
+        When IPython is started normally, both parameters are None: a new module
+        is created automatically, and its __dict__ used as the namespace.
+        
+        If only user_module is provided, its __dict__ is used as the namespace.
+        If only user_ns is provided, a dummy module is created, and user_ns
+        becomes the global namespace. If both are provided (as they may be
+        when embedding), user_ns is the local namespace, and user_module
+        provides the global namespace.
 
         Parameters
         ----------
         user_module : module, optional
             The current user module in which IPython is being run. If None,
             a clean module will be created.
+        user_ns : dict, optional
+            A namespace in which to run interactive commands.
 
         Returns
         -------
-        A module object.
+        A tuple of user_module and user_ns, each properly initialised.
         """
+        if user_module is None and user_ns is not None:
+            user_ns.setdefault("__name__", "__main__")
+            class DummyMod(object):
+                "A dummy module used for IPython's interactive namespace."
+                pass
+            user_module = DummyMod()
+            user_module.__dict__ = user_ns
+            
         if user_module is None:
             user_module = types.ModuleType("__main__",
                 doc="Automatically created module for IPython interactive environment")
@@ -966,8 +978,11 @@ class InteractiveShell(SingletonConfigurable, Magic):
         # http://mail.python.org/pipermail/python-dev/2001-April/014068.html
         user_module.__dict__.setdefault('__builtin__',__builtin__)
         user_module.__dict__.setdefault('__builtins__',__builtin__)
+        
+        if user_ns is None:
+            user_ns = user_module.__dict__
 
-        return user_module
+        return user_module, user_ns
 
     def init_sys_modules(self):
         # We need to insert into sys.modules something that looks like a
@@ -1075,13 +1090,16 @@ class InteractiveShell(SingletonConfigurable, Magic):
         # The main execution namespaces must be cleared very carefully,
         # skipping the deletion of the builtin-related keys, because doing so
         # would cause errors in many object's __del__ methods.
-        for ns in [self.user_ns, self.user_global_ns]:
-            drop_keys = set(ns.keys())
-            drop_keys.discard('__builtin__')
-            drop_keys.discard('__builtins__')
-            for k in drop_keys:
-                del ns[k]
-
+        if self.user_ns is not self.user_global_ns:
+            self.user_ns.clear()
+        ns = self.user_global_ns
+        drop_keys = set(ns.keys())
+        drop_keys.discard('__builtin__')
+        drop_keys.discard('__builtins__')
+        drop_keys.discard('__name__')
+        for k in drop_keys:
+            del ns[k]
+        
         # Restore the user namespaces to minimal usability
         self.init_user_ns()
 
diff --git a/IPython/core/tests/test_interactiveshell.py b/IPython/core/tests/test_interactiveshell.py
index 6ed510f..30496f6 100644
--- a/IPython/core/tests/test_interactiveshell.py
+++ b/IPython/core/tests/test_interactiveshell.py
@@ -25,6 +25,7 @@ import shutil
 import tempfile
 import unittest
 from os.path import join
+import sys
 from StringIO import StringIO
 
 from IPython.testing import decorators as dec
@@ -128,6 +129,7 @@ class InteractiveShellTestCase(unittest.TestCase):
         f = IPython.core.formatters.PlainTextFormatter()
         f([Spam(),Spam()])
     
+
     def test_future_flags(self):
         """Check that future flags are used for parsing code (gh-777)"""
         ip = get_ipython()
@@ -151,6 +153,37 @@ class InteractiveShellTestCase(unittest.TestCase):
         finally:
             # Reset compiler flags so we don't mess up other tests.
             ip.compile.reset_compiler_flags()
+    
+    def test_can_pickle(self):
+        "Can we pickle objects defined interactively (GH-29)"
+        ip = get_ipython()
+        ip.reset()
+        ip.run_cell(("class Mylist(list):\n"
+                     "    def __init__(self,x=[]):\n"
+                     "        list.__init__(self,x)"))
+        ip.run_cell("w=Mylist([1,2,3])")
+        
+        from cPickle import dumps
+        
+        # We need to swap in our main module - this is only necessary
+        # inside the test framework, because IPython puts the interactive module
+        # in place (but the test framework undoes this).
+        _main = sys.modules['__main__']
+        sys.modules['__main__'] = ip.user_module
+        try:
+            res = dumps(ip.user_ns["w"])
+        finally:
+            sys.modules['__main__'] = _main
+        self.assertTrue(isinstance(res, bytes))
+        
+    def test_global_ns(self):
+        "Code in functions must be able to access variables outside them."
+        ip = get_ipython()
+        ip.run_cell("a = 10")
+        ip.run_cell(("def f(x):"
+                     "    return x + a"))
+        ip.run_cell("b = f(12)")
+        self.assertEqual(ip.user_ns["b"], 22)
 
     def test_bad_custom_tb(self):
         """Check that InteractiveShell is protected from bad custom exception handlers"""
diff --git a/IPython/testing/globalipapp.py b/IPython/testing/globalipapp.py
index 3544f67..44717ec 100644
--- a/IPython/testing/globalipapp.py
+++ b/IPython/testing/globalipapp.py
@@ -205,8 +205,7 @@ def start_ipython():
     # can capture subcommands and print them to Python's stdout, otherwise the
     # doctest machinery would miss them.
     shell.system = py3compat.MethodType(xsys, shell)
-
-
+    
     shell._showtraceback = py3compat.MethodType(_showtraceback, shell)
 
     # IPython is ready, now clean up some global state...
diff --git a/IPython/testing/plugin/ipdoctest.py b/IPython/testing/plugin/ipdoctest.py
index 6948aa3..50cc72b 100644
--- a/IPython/testing/plugin/ipdoctest.py
+++ b/IPython/testing/plugin/ipdoctest.py
@@ -271,6 +271,8 @@ class DocTestCase(doctests.DocTestCase):
             # for IPython examples *only*, we swap the globals with the ipython
             # namespace, after updating it with the globals (which doctest
             # fills with the necessary info from the module being tested).
+            self.user_ns_orig = {}
+            self.user_ns_orig.update(_ip.user_ns)
             _ip.user_ns.update(self._dt_test.globs)
             self._dt_test.globs = _ip.user_ns
             # IPython must protect the _ key in the namespace (it can't exist)
@@ -286,6 +288,8 @@ class DocTestCase(doctests.DocTestCase):
         # teardown doesn't destroy the ipython namespace
         if isinstance(self._dt_test.examples[0],IPExample):
             self._dt_test.globs = self._dt_test_globs_ori
+            _ip.user_ns.clear()
+            _ip.user_ns.update(self.user_ns_orig)
             # Restore the behavior of the '_' key in the user namespace to
             # normal after each doctest, so that unittests behave normally
             _ip.user_ns.protect_underscore = False