From 01ea6c6b7188744c064663f124c908ef4313b9dc 2013-08-27 21:08:34
From: MinRK <benjaminrk@gmail.com>
Date: 2013-08-27 21:08:34
Subject: [PATCH] adjust strip_prompts logic a bit further

- initial-only prompt is only stripped from the first line (only change from master)
- any prompt is stripped from the second line
- if a prompt is found in the first two lines,
  strip from the rest of the block
---

diff --git a/IPython/core/inputtransformer.py b/IPython/core/inputtransformer.py
index 47ebda3..84280a2 100644
--- a/IPython/core/inputtransformer.py
+++ b/IPython/core/inputtransformer.py
@@ -359,17 +359,26 @@ def cellmagic(end_on_blank_line=False):
         line = tpl % (magic_name, first, u'\n'.join(body))
 
 
-def _strip_prompts(prompt_re, continuation_re=None):
+def _strip_prompts(prompt_re, initial_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.
+    Parameters
+    ----------
+    prompt_re : regular expression
+        A regular expression matching any input prompt (including continuation)
+    initial_re : regular expression, optional
+        A regular expression matching only the initial prompt, but not continuation.
+        If no initial expression is given, prompt_re will be used everywhere.
+        Used mainly for plain Python prompts, where the continuation prompt
+        ``...`` is a valid Python expression in Python 3, so shouldn't be stripped.
     
-    If continuation_re is unspecified, prompt_re will be used for both.
+    If initial_re and prompt_re differ,
+    only initial_re will be tested against the first line.
+    If any prompt is found on the first two lines,
+    prompts will be stripped from the rest of the block.
     """
-    if continuation_re is None:
-        continuation_re = prompt_re
+    if initial_re is None:
+        initial_re = prompt_re
     line = ''
     while True:
         line = (yield line)
@@ -377,26 +386,22 @@ def _strip_prompts(prompt_re, continuation_re=None):
         # First line of cell
         if line is None:
             continue
-        out, n1 = prompt_re.subn('', line, count=1)
+        out, n1 = initial_re.subn('', line, count=1)
         line = (yield out)
         
-        # Second line of cell, because people often copy from just after the
-        # first prompt, so we might not see it in the first line.
         if line is None:
             continue
-        # 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)
+        # check for any prompt on the second line of the cell,
+        # because people often copy from just after the first prompt,
+        # so we might not see it in the first line.
+        out, n2 = prompt_re.subn('', line, count=1)
         line = (yield out)
         
         if n1 or n2:
-            # Found the input prompt in the first two lines - check for it in
+            # Found a prompt in the first two lines - check for it in
             # the rest of the cell as well.
             while line is not None:
-                line = (yield continuation_re.sub('', line, count=1))
+                line = (yield prompt_re.sub('', line, count=1))
         
         else:
             # Prompts not in input - wait for reset
@@ -407,18 +412,17 @@ def _strip_prompts(prompt_re, continuation_re=None):
 def classic_prompt():
     """Strip the >>>/... prompts of the Python interactive shell."""
     # FIXME: non-capturing version (?:...) usable?
-    prompt_re = re.compile(r'^(>>> ?)')
-    continuation_re = re.compile(r'^(>>> ?|\.\.\. ?)')
-    return _strip_prompts(prompt_re, continuation_re)
+    prompt_re = re.compile(r'^(>>> ?|\.\.\. ?)')
+    initial_re = re.compile(r'^(>>> ?)')
+    return _strip_prompts(prompt_re, initial_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+\]: )')
-    continuation_re = re.compile(r'^(In \[\d+\]: |\ \ \ \.\.\.+: )')
-    return _strip_prompts(prompt_re, continuation_re)
+    prompt_re = re.compile(r'^(In \[\d+\]: |\ \ \ \.\.\.+: )')
+    return _strip_prompts(prompt_re)
 
 
 @CoroutineInputTransformer.wrap
diff --git a/IPython/core/tests/test_inputtransformer.py b/IPython/core/tests/test_inputtransformer.py
index 39a0df3..ab585e0 100644
--- a/IPython/core/tests/test_inputtransformer.py
+++ b/IPython/core/tests/test_inputtransformer.py
@@ -179,8 +179,12 @@ syntax_ml = \
           ('... 123"""','123"""'),
           ],
          [('a="""','a="""'),
-          ('... 123','... 123'),
-          ('... 456"""','... 456"""'),
+          ('... 123','123'),
+          ('... 456"""','456"""'),
+          ],
+         [('a="""','a="""'),
+          ('>>> 123','123'),
+          ('... 456"""','456"""'),
           ],
          [('a="""','a="""'),
           ('123','123'),
@@ -189,7 +193,7 @@ syntax_ml = \
          [('....__class__','....__class__'),
          ],
          [('a=5', 'a=5'),
-          ('...', '...'),
+          ('...', ''),
          ],
          [('>>> def f(x):', 'def f(x):'),
           ('...', ''),
@@ -206,8 +210,12 @@ syntax_ml = \
           ('   ...: 123"""','123"""'),
           ],
          [('a="""','a="""'),
-          ('   ...: 123','   ...: 123'),
-          ('   ...: 456"""','   ...: 456"""'),
+          ('   ...: 123','123'),
+          ('   ...: 456"""','456"""'),
+          ],
+         [('a="""','a="""'),
+          ('In [1]: 123','123'),
+          ('   ...: 456"""','456"""'),
           ],
          [('a="""','a="""'),
           ('123','123'),