From 8822e910b98091f43872801a90fee8c97d821545 2011-12-14 01:44:29
From: MinRK <benjaminrk@gmail.com>
Date: 2011-12-14 01:44:29
Subject: [PATCH] allow access to user_ns in prompt_manager

adds UserNSFormatter for falling back on shell.user_ns / builtins for
name resolution.

Closes #1151, rebased to prevent recursive merge.

---

diff --git a/IPython/core/prompts.py b/IPython/core/prompts.py
index 28cbee5..ca276ed 100644
--- a/IPython/core/prompts.py
+++ b/IPython/core/prompts.py
@@ -26,6 +26,8 @@ import socket
 import sys
 import time
 
+from string import Formatter
+
 from IPython.config.configurable import Configurable
 from IPython.core import release
 from IPython.utils import coloransi
@@ -241,6 +243,26 @@ def _lenlastline(s):
         return 0
     return len(s.splitlines()[-1])
 
+
+class UserNSFormatter(Formatter):
+    """A Formatter that falls back on a shell's user_ns and __builtins__ for name resolution"""
+    def __init__(self, shell):
+        self.shell = shell
+
+    def get_value(self, key, args, kwargs):
+        # try regular formatting first:
+        try:
+            return Formatter.get_value(self, key, args, kwargs)
+        except Exception:
+            pass
+        # next, look in user_ns and builtins:
+        for container in (self.shell.user_ns, __builtins__):
+            if key in container:
+                return container[key]
+        # nothing found, put error message in its place
+        return "<ERROR: '%s' not found>" % key
+
+
 class PromptManager(Configurable):
     """This is the primary interface for producing IPython's prompts."""
     shell = Instance('IPython.core.interactiveshell.InteractiveShellABC')
@@ -293,6 +315,7 @@ class PromptManager(Configurable):
         self.color_scheme_table = coloransi.ColorSchemeTable([PColNoColors,
                                     PColLinux, PColLightBG], self.color_scheme)
         
+        self._formatter = UserNSFormatter(shell)
         # Prepare templates & numbers of invisible characters
         self.update_prompt('in', self.in_template)
         self.update_prompt('in2', self.in2_template)
@@ -357,7 +380,7 @@ class PromptManager(Configurable):
         prompt = colors.prompt + self.templates[name] + colors.normal
         
         # Fill in required fields
-        return prompt.format(**fmtargs)
+        return self._formatter.format(prompt, **fmtargs)
     
     def _render_rewrite(self, color=True):
         """Render the ---> rewrite prompt."""
diff --git a/IPython/core/tests/test_prompts.py b/IPython/core/tests/test_prompts.py
index d3613b3..22ab04a 100644
--- a/IPython/core/tests/test_prompts.py
+++ b/IPython/core/tests/test_prompts.py
@@ -38,6 +38,25 @@ class PromptTests(unittest.TestCase):
     
         tt.check_pairs(do_translate, pairs)
     
+    def test_user_ns(self):
+        self.pm.color_scheme = 'NoColor'
+        ip.ex("foo='bar'")
+        self.pm.in_template = "In [{foo}]"
+        prompt = self.pm.render('in')
+        self.assertEquals(prompt, u'In [bar]')
+
+    def test_builtins(self):
+        self.pm.color_scheme = 'NoColor'
+        self.pm.in_template = "In [{int}]"
+        prompt = self.pm.render('in')
+        self.assertEquals(prompt, u"In [<type 'int'>]")
+
+    def test_undefined(self):
+        self.pm.color_scheme = 'NoColor'
+        self.pm.in_template = "In [{foo_dne}]"
+        prompt = self.pm.render('in')
+        self.assertEquals(prompt, u"In [<ERROR: 'foo_dne' not found>]")
+
     def test_render(self):
         self.pm.in_template = r'\#>'
         self.assertEqual(self.pm.render('in',color=False), '%d>' % ip.execution_count)