gh-95191: IDLE: Include prompts when saving Shell #95554

1 files changed, 19 insertions, 16 deletions

diff --git a/Lib/idlelib/pyshell.py b/Lib/idlelib/pyshell.py
index 38061229622..e68233a5a41 100755
--- a/Lib/idlelib/pyshell.py
+++ b/Lib/idlelib/pyshell.py
@@ -986,6 +986,23 @@ class PyShell(OutputWindow):
     def get_standard_extension_names(self):
         return idleConf.GetExtensions(shell_only=True)
 
+    def get_prompt_text(self, first, last):
+        """Return text between first and last with prompts added."""
+        text = self.text.get(first, last)
+        lineno_range = range(
+            int(float(first)),
+            int(float(last))
+         )
+        prompts = [
+            self.shell_sidebar.line_prompts.get(lineno)
+            for lineno in lineno_range
+        ]
+        return "\n".join(
+            line if prompt is None else f"{prompt} {line}"
+            for prompt, line in zip(prompts, text.splitlines())
+        ) + "\n"
+
+
     def copy_with_prompts_callback(self, event=None):
         """Copy selected lines to the clipboard, with prompts.
 
@@ -1002,23 +1019,9 @@ class PyShell(OutputWindow):
         sellast = text.index('sel.last')
         if sellast[-1] != '0':
             sellast = text.index("sel.last+1line linestart")
-
-        selected_text = self.text.get(selfirst, sellast)
-        selection_lineno_range = range(
-            int(float(selfirst)),
-            int(float(sellast))
-        )
-        prompts = [
-            self.shell_sidebar.line_prompts.get(lineno)
-            for lineno in selection_lineno_range
-        ]
-        selected_text_with_prompts = "\n".join(
-            line if prompt is None else f"{prompt} {line}"
-            for prompt, line in zip(prompts, selected_text.splitlines())
-        ) + "\n"
-
         text.clipboard_clear()
-        text.clipboard_append(selected_text_with_prompts)
+        prompt_text = self.get_prompt_text(selfirst, sellast)
+        text.clipboard_append(prompt_text)
 
     reading = False
     executing = False