From f549a21f3fa831bd3f7630bfab9912012ef4ff46 2013-03-31 09:03:54
From: Thomas Kluyver <takowl@gmail.com>
Date: 2013-03-31 09:03:54
Subject: [PATCH] Add TokenInputTransformer

---

diff --git a/IPython/core/inputtransformer.py b/IPython/core/inputtransformer.py
index b016e30..4440b83 100644
--- a/IPython/core/inputtransformer.py
+++ b/IPython/core/inputtransformer.py
@@ -110,6 +110,50 @@ class CoroutineInputTransformer(InputTransformer):
         """
         return self.coro.send(None)
 
+class TokenInputTransformer(InputTransformer):
+    """Wrapper for a token-based input transformer.
+    
+    func should accept a list of tokens (5-tuples, see tokenize docs), and
+    return an iterable which can be passed to tokenize.untokenize().
+    """
+    def __init__(self, func):
+        self.func = func
+        self.current_line = ""
+        self.tokenizer = tokenize.generate_tokens(self.get_line)
+        self.line_used= False
+    
+    def get_line(self):
+        if self.line_used:
+            raise tokenize.TokenError
+        self.line_used = True
+        return self.current_line
+    
+    def push(self, line):
+        self.current_line += line + "\n"
+        self.line_used = False
+        tokens = []
+        try:
+            for intok in self.tokenizer:
+                tokens.append(intok)
+                if intok[0] in (tokenize.NEWLINE, tokenize.NL):
+                    # Stop before we try to pull a line we don't have yet
+                    break
+        except tokenize.TokenError:
+            # Multi-line statement - stop and try again with the next line
+            self.tokenizer = tokenize.generate_tokens(self.get_line)
+            return None
+        
+        self.current_line = ""
+        # Python bug 8478 - untokenize doesn't work quite correctly with a
+        # generator. We call list() to avoid this.
+        return tokenize.untokenize(list(self.func(tokens))).rstrip('\n')
+    
+    def reset(self):
+        l = self.current_line
+        self.current_line = ""
+        if l:
+            return l.rstrip('\n')
+
 
 # Utilities
 def _make_help_call(target, esc, lspace, next_input=None):
diff --git a/IPython/core/tests/test_inputtransformer.py b/IPython/core/tests/test_inputtransformer.py
index 4ed0551..5e9df92 100644
--- a/IPython/core/tests/test_inputtransformer.py
+++ b/IPython/core/tests/test_inputtransformer.py
@@ -1,3 +1,4 @@
+import tokenize
 import unittest
 import nose.tools as nt
 
@@ -324,3 +325,44 @@ def test_has_comment():
              ('a #comment not "string"', True),
              ]
     tt.check_pairs(ipt.has_comment, tests)
+
+@ipt.TokenInputTransformer.wrap
+def decistmt(tokens):
+    """Substitute Decimals for floats in a string of statements.
+
+    Based on an example from the tokenize module docs.
+    """
+    result = []
+    for toknum, tokval, _, _, _  in tokens:
+        if toknum == tokenize.NUMBER and '.' in tokval:  # replace NUMBER tokens
+            for newtok in [
+                (tokenize.NAME, 'Decimal'),
+                (tokenize.OP, '('),
+                (tokenize.STRING, repr(tokval)),
+                (tokenize.OP, ')')
+            ]:
+                yield newtok
+        else:
+            yield (toknum, tokval)
+
+
+
+def test_token_input_transformer():
+    tests = [(u'1.2', u_fmt(u"Decimal ({u}'1.2')")),
+             (u'"1.2"', u'"1.2"'),
+             ]
+    tt.check_pairs(transform_and_reset(decistmt), tests)
+    ml_tests = \
+    [ [(u"a = 1.2; b = '''x", None),
+       (u"y'''", u_fmt(u"a =Decimal ({u}'1.2');b ='''x\ny'''")),
+      ],
+      [(u"a = [1.2,", u_fmt(u"a =[Decimal ({u}'1.2'),")),
+       (u"3]", u"3 ]"),
+      ],
+      [(u"a = '''foo", None),  # Test resetting when within a multi-line string
+       (u"bar", None),
+       (None, u"a = '''foo\nbar"),
+      ],
+    ]
+    for example in ml_tests:
+        transform_checker(example, decistmt)