Override get_session_formatted_for_file in curtsies Repl

sebastinas · sebastinas · commit bae3fd05a44d · 2021-05-16T22:56:34.000+02:00
diff --git a/bpython/curtsiesfrontend/repl.py b/bpython/curtsiesfrontend/repl.py
@@ -10,6 +10,7 @@
 import tempfile
 import time
 import unicodedata
+from enum import Enum
 
 import blessings
 import greenlet
@@ -50,7 +51,6 @@
 from ..repl import (
     Repl,
     SourceNotFound,
-    LineTypeTranslator as LineType,
 )
 from ..translations import _
 
@@ -70,6 +70,14 @@
 MAX_EVENTS_POSSIBLY_NOT_PASTE = 20
 
 
+class LineType(Enum):
+    """Used when adding a tuple to all_logical_lines, to get input / output values
+    having to actually type/know the strings"""
+
+    INPUT = "input"
+    OUTPUT = "output"
+
+
 class FakeStdin:
     """The stdin object user code will reference
 
@@ -365,7 +373,7 @@ def __init__(
         # Entries are tuples, where
         #   - the first element the line (string, not fmtsr)
         #   - the second element is one of 2 global constants: "input" or "output"
-        #     (use LineType.INPUT or LineType.OUTPUT to avoid typing these strings)
+        #     (use LineTypeTranslator.INPUT or LineTypeTranslator.OUTPUT to avoid typing these strings)
         self.all_logical_lines = []
 
         # formatted version of lines in the buffer kept around so we can
@@ -2032,6 +2040,17 @@ def key_help_text(self):
             "{} : {}".format(func.rjust(max_func), key) for func, key in pairs
         )
 
+    def get_session_formatted_for_file(self) -> str:
+        def process():
+            for line, lineType in self.all_logical_lines:
+                if lineType == LineType.INPUT:
+                    yield line
+                elif line.rstrip():
+                    yield "# OUT: %s" % line
+            yield "### %s" % self.current_line
+
+        return "\n".join(process())
+
 
 def is_nop(char):
     return unicodedata.category(str(char)) == "Cc"
diff --git a/bpython/repl.py b/bpython/repl.py
@@ -33,7 +33,6 @@
 import textwrap
 import time
 import traceback
-from enum import Enum
 from itertools import takewhile
 from pathlib import Path
 from pygments.lexers import Python3Lexer
@@ -355,14 +354,6 @@ class SourceNotFound(Exception):
     """Exception raised when the requested source could not be found."""
 
 
-class LineTypeTranslator(Enum):
-    """Used when adding a tuple to all_logical_lines, to get input / output values
-    having to actually type/know the strings"""
-
-    INPUT = "input"
-    OUTPUT = "output"
-
-
 class Repl:
     """Implements the necessary guff for a Python-repl-alike interface
 
@@ -778,37 +769,23 @@ def line_is_empty(line):
             indentation = 0
         return indentation
 
-    def get_session_formatted_for_file(self):
+    def get_session_formatted_for_file(self) -> str:
         """Format the stdout buffer to something suitable for writing to disk,
         i.e. without >>> and ... at input lines and with "# OUT: " prepended to
         output lines and "### " prepended to current line"""
 
-        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()
+        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 f"# OUT: {line}"
+        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 f"# OUT: {line}"
 
-            return "\n".join(process())
+        return "\n".join(process())
 
     def write2file(self):
         """Prompt for a filename and write the current contents of the stdout
