From 250c2bd94d63322e583e0dbfbf41b14f129d320d 2013-08-19 05:53:36
From: MinRK <benjaminrk@gmail.com>
Date: 2013-08-19 05:53:36
Subject: [PATCH] only strip continuation prompts if regular prompts seen first

should be discussed before merge (ping @takluyver),
as it is a significant and sensitive change.

I'm also not sure that it's the right fix, but worth a try.

closes #4059
---

diff --git a/IPython/core/inputtransformer.py b/IPython/core/inputtransformer.py
index 1a5f1bb..47ebda3 100644
--- a/IPython/core/inputtransformer.py
+++ b/IPython/core/inputtransformer.py
@@ -359,8 +359,17 @@ def cellmagic(end_on_blank_line=False):
         line = tpl % (magic_name, first, u'\n'.join(body))
 
 
-def _strip_prompts(prompt_re):
-    """Remove matching input prompts from a block of input."""
+def _strip_prompts(prompt_re, continuation_re=None):
+    """Remove matching input prompts from a block of input.
+    
+    prompt_re is stripped only once, on either the first or second line.
+    If prompt_re is found on one of the first two lines,
+    continuation_re is stripped from lines thereafter.
+    
+    If continuation_re is unspecified, prompt_re will be used for both.
+    """
+    if continuation_re is None:
+        continuation_re = prompt_re
     line = ''
     while True:
         line = (yield line)
@@ -375,14 +384,19 @@ def _strip_prompts(prompt_re):
         # first prompt, so we might not see it in the first line.
         if line is None:
             continue
-        out, n2 = prompt_re.subn('', line, count=1)
+        # check for first prompt if not found on first line, continuation otherwise
+        if n1:
+            pat = continuation_re
+        else:
+            pat = prompt_re
+        out, n2 = pat.subn('', line, count=1)
         line = (yield out)
         
         if n1 or n2:
             # Found the input prompt in the first two lines - check for it in
             # the rest of the cell as well.
             while line is not None:
-                line = (yield prompt_re.sub('', line, count=1))
+                line = (yield continuation_re.sub('', line, count=1))
         
         else:
             # Prompts not in input - wait for reset
@@ -393,16 +407,18 @@ def _strip_prompts(prompt_re):
 def classic_prompt():
     """Strip the >>>/... prompts of the Python interactive shell."""
     # FIXME: non-capturing version (?:...) usable?
-    prompt_re = re.compile(r'^(>>> ?|\.\.\. ?)')
-    return _strip_prompts(prompt_re)
+    prompt_re = re.compile(r'^(>>> ?)')
+    continuation_re = re.compile(r'^(>>> ?|\.\.\. ?)')
+    return _strip_prompts(prompt_re, continuation_re)
 
 @CoroutineInputTransformer.wrap
 def ipy_prompt():
     """Strip IPython's In [1]:/...: prompts."""
     # FIXME: non-capturing version (?:...) usable?
     # FIXME: r'^(In \[\d+\]: | {3}\.{3,}: )' clearer?
-    prompt_re = re.compile(r'^(In \[\d+\]: |\ \ \ \.\.\.+: )')
-    return _strip_prompts(prompt_re)
+    prompt_re = re.compile(r'^(In \[\d+\]: )')
+    continuation_re = re.compile(r'^(In \[\d+\]: |\ \ \ \.\.\.+: )')
+    return _strip_prompts(prompt_re, continuation_re)
 
 
 @CoroutineInputTransformer.wrap
diff --git a/IPython/core/tests/test_inputtransformer.py b/IPython/core/tests/test_inputtransformer.py
index 56d4c55..39a0df3 100644
--- a/IPython/core/tests/test_inputtransformer.py
+++ b/IPython/core/tests/test_inputtransformer.py
@@ -179,13 +179,18 @@ syntax_ml = \
           ('... 123"""','123"""'),
           ],
          [('a="""','a="""'),
-          ('... 123','123'),
-          ('... 456"""','456"""'),
+          ('... 123','... 123'),
+          ('... 456"""','... 456"""'),
           ],
          [('a="""','a="""'),
           ('123','123'),
           ('... 456"""','... 456"""'),
           ],
+         [('....__class__','....__class__'),
+         ],
+         [('a=5', 'a=5'),
+          ('...', '...'),
+         ],
          [('>>> def f(x):', 'def f(x):'),
           ('...', ''),
           ('...     return x', '    return x'),
@@ -201,8 +206,8 @@ syntax_ml = \
           ('   ...: 123"""','123"""'),
           ],
          [('a="""','a="""'),
-          ('   ...: 123','123'),
-          ('   ...: 456"""','456"""'),
+          ('   ...: 123','   ...: 123'),
+          ('   ...: 456"""','   ...: 456"""'),
           ],
          [('a="""','a="""'),
           ('123','123'),