Restore original save behavior for cli and urwid

thomasballinger · sebastinas · commit cd01afdd9ba2 · 2020-10-12T21:52:14.000+02:00
diff --git a/bpython/curtsiesfrontend/repl.py b/bpython/curtsiesfrontend/repl.py
@@ -1164,6 +1164,8 @@ def push(self, line, insert_into_history=True):
 
         If the interpreter successfully runs the code, clear the buffer
         """
+        # Note that push() overrides its parent without calling it, unlike
+        # urwid and cli which implement custom behavior and call repl.Repl.push
         if self.paste_mode:
             self.saved_indent = 0
         else:
diff --git a/bpython/repl.py b/bpython/repl.py
@@ -832,15 +832,31 @@ def get_session_formatted_for_file(self):
         i.e. without >>> and ... at input lines and with "# OUT: " prepended to
         output lines and "### " prepended to current line"""
 
-        def process():
-            for line, lineType in self.all_logical_lines:
-                if lineType == LineTypeTranslator.INPUT:
-                    yield line
-                elif line.rstrip():
-                    yield "# OUT: %s" % line
-            yield "### %s" % self.current_line
-
-        return "\n".join(process())
+        if hasattr(self, 'all_logical_lines'):
+            # Curtsies
+
+            def process():
+                for line, lineType in self.all_logical_lines:
+                    if lineType == LineTypeTranslator.INPUT:
+                        yield line
+                    elif line.rstrip():
+                        yield "# OUT: %s" % line
+                yield "### %s" % self.current_line
+
+            return "\n".join(process())
+
+        else:  # cli and Urwid
+            session_output = self.getstdout()
+
+            def process():
+                for line in session_output.split("\n"):
+                    if line.startswith(self.ps1):	
+                        yield line[len(self.ps1) :]	
+                    elif line.startswith(self.ps2):	
+                        yield line[len(self.ps2) :]	
+                    elif line.rstrip():	        
+                        yield "# OUT: %s" % (line,)
+            return '\n'.join(process())
 
     def write2file(self):
         """Prompt for a filename and write the current contents of the stdout
@@ -952,6 +968,7 @@ def do_pastebin(self, s):
     def push(self, s, insert_into_history=True):
         """Push a line of code onto the buffer so it can process it all
         at once when a code block ends"""
+        # This push method is used by cli and urwid, but not curtsies
         s = s.rstrip("\n")
         self.buffer.append(s)
 
