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

4 files changed, 42 insertions, 24 deletions

diff --git a/Lib/idlelib/NEWS.txt b/Lib/idlelib/NEWS.txt
index 0e3a50b..ce95e2f 100644
--- a/Lib/idlelib/NEWS.txt
+++ b/Lib/idlelib/NEWS.txt
@@ -4,6 +4,8 @@ Released on 2022-10-03
 =========================
 
 
+gh-95191: Include prompts when saving Shell (interactive input/output).
+
 gh-95511: Fix the Shell context menu copy-with-prompts bug of copying
 an extra line when one selects whole lines.
 
diff --git a/Lib/idlelib/idle_test/test_iomenu.py b/Lib/idlelib/idle_test/test_iomenu.py
index e338893..2fb836d 100644
--- a/Lib/idlelib/idle_test/test_iomenu.py
+++ b/Lib/idlelib/idle_test/test_iomenu.py
@@ -1,10 +1,12 @@
 "Test , coverage 17%."
 
-from idlelib import iomenu, util
+from idlelib import iomenu
 import unittest
 from test.support import requires
 from tkinter import Tk
 from idlelib.editor import EditorWindow
+from idlelib import util
+from idlelib.idle_test.mock_idle import Func
 
 
 class IOBindingTest(unittest.TestCase):
@@ -36,9 +38,14 @@ class IOBindingTest(unittest.TestCase):
         io = self.io
         fix = io.fixnewlines
         text = io.editwin.text
+
+        # Make the editor temporarily look like Shell.
         self.editwin.interp = None
-        eq(fix(), '')
-        del self.editwin.interp
+        shelltext = '>>> if 1'
+        self.editwin.get_prompt_text = Func(result=shelltext)
+        eq(fix(), shelltext)  # Get... call and '\n' not added.
+        del self.editwin.interp, self.editwin.get_prompt_text
+
         text.insert(1.0, 'a')
         eq(fix(), 'a'+io.eol_convention)
         eq(text.get('1.0', 'end-1c'), 'a\n')
diff --git a/Lib/idlelib/iomenu.py b/Lib/idlelib/iomenu.py
index 327e885..af8159c 100644
--- a/Lib/idlelib/iomenu.py
+++ b/Lib/idlelib/iomenu.py
@@ -251,11 +251,17 @@ class IOBinding:
             return False
 
     def fixnewlines(self):
-        "Return text with final \n if needed and os eols."
-        if (self.text.get("end-2c") != '\n'
-            and not hasattr(self.editwin, "interp")):  # Not shell.
-            self.text.insert("end-1c", "\n")
-        text = self.text.get("1.0", "end-1c")
+        """Return text with os eols.
+
+        Add prompts if shell else final \n if missing.
+        """
+
+        if hasattr(self.editwin, "interp"):  # Saving shell.
+            text = self.editwin.get_prompt_text('1.0', self.text.index('end-1c'))
+        else:
+            if self.text.get("end-2c") != '\n':
+                self.text.insert("end-1c", "\n")  # Changes 'end-1c' value.
+            text = self.text.get('1.0', "end-1c")
         if self.eol_convention != "\n":
             text = text.replace("\n", self.eol_convention)
         return text
diff --git a/Lib/idlelib/pyshell.py b/Lib/idlelib/pyshell.py
index 3806122..e68233a 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