don't hardcode prompts

thomasballinger · thomasballinger · commit b2a3bf202f25 · 2014-08-06T11:03:23.000-04:00
diff --git a/bpython/curtsiesfrontend/repl.py b/bpython/curtsiesfrontend/repl.py
@@ -233,7 +233,7 @@ def smarter_request_refresh():
         self.interact = self.status_bar # overwriting what bpython.Repl put there
                                         # interact is called to interact with the status bar,
                                         # so we're just using the same object
-        self.current_line = '' # line currently being edited, without '>>> '
+        self.current_line = '' # line currently being edited, without ps1 (usually ">>> ")
         self.current_stdouterr_line = '' # current line of output - stdout and stdin go here
         self.display_lines = [] # lines separated whenever logical line
                                 # length goes over what the terminal width
@@ -524,8 +524,10 @@ def send_current_block_to_external_editor(self, filename=None):
 
     def send_session_to_external_editor(self, filename=None):
         for_editor = '### current bpython session - file will be reevaluated, ### lines will not be run\n'.encode('utf8')
-        for_editor += ('\n'.join(line[4:] if line[:4] in ('... ', '>>> ') else '### '+line
-                       for line in self.getstdout().split('\n')).encode('utf8'))
+        for_editor += ('\n'.join(line[len(self.ps1):] if line.startswith(self.ps1) else
+                                 (line[len(self.ps2):] if line.startswith(self.ps2) else
+                                 '### '+line)
+                                 for line in self.getstdout().split('\n')).encode('utf8'))
         text = self.send_to_external_editor(for_editor)
         lines = text.split('\n')
         self.history = [line for line in lines if line[:4] != '### ']
@@ -549,7 +551,7 @@ def add_normal_character(self, char):
                              char +
                              self.current_line[self.cursor_offset:])
         self.cursor_offset += 1
-        if self.config.cli_trim_prompts and self.current_line.startswith(">>> "):
+        if self.config.cli_trim_prompts and self.current_line.startswith(self.ps1):
             self.current_line = self.current_line[4:]
             self.cursor_offset = max(0, self.cursor_offset - 4)
 
