load extensions in builtin trap

minrk · minrk · commit 516dcafe9e52 · 2013-05-06T13:44:25.000-07:00
ensures `get_ipython` et al. are available.
diff --git a/IPython/core/extensions.py b/IPython/core/extensions.py
@@ -92,15 +92,16 @@ def load_extension(self, module_str):
             return "already loaded"
         
         from IPython.utils.syspathcontext import prepended_to_syspath
-
-        if module_str not in sys.modules:
-            with prepended_to_syspath(self.ipython_extension_dir):
-                __import__(module_str)
-        mod = sys.modules[module_str]
-        if self._call_load_ipython_extension(mod):
-            self.loaded.add(module_str)
-        else:
-            return "no load function"
+        
+        with self.shell.builtin_trap:
+            if module_str not in sys.modules:
+                with prepended_to_syspath(self.ipython_extension_dir):
+                    __import__(module_str)
+            mod = sys.modules[module_str]
+            if self._call_load_ipython_extension(mod):
+                self.loaded.add(module_str)
+            else:
+                return "no load function"
 
     def unload_extension(self, module_str):
         """Unload an IPython extension by its module name.
diff --git a/IPython/core/tests/test_extension.py b/IPython/core/tests/test_extension.py
@@ -19,6 +19,12 @@ def load_ipython_extension(ip):
     print("Running ext2 load")
 """
 
+ext3_content = """
+def load_ipython_extension(ip):
+    ip2 = get_ipython()
+    print(ip is ip2)
+"""
+
 def test_extension_loading():
     em = get_ipython().extension_manager
     with TemporaryDirectory() as td:
@@ -68,6 +74,23 @@ def test_extension_loading():
             with tt.AssertPrints("Running ext2 load"):
                 em.reload_extension('ext2')
 
+
+def test_extension_builtins():
+    em = get_ipython().extension_manager
+    with TemporaryDirectory() as td:
+        ext3 = os.path.join(td, 'ext3.py')
+        with open(ext3, 'w') as f:
+            f.write(ext3_content)
+        
+        assert 'ext3' not in em.loaded
+        
+        with prepended_to_syspath(td):
+            # Load extension
+            with tt.AssertPrints("True"):
+                assert em.load_extension('ext3') is None
+            assert 'ext3' in em.loaded
+
+
 def test_non_extension():
     em = get_ipython().extension_manager
     nt.assert_equal(em.load_extension('sys'), "no load function")
