Fix handling of arg prompt and add tests for it

- and for `unrecoverable` exceptions

Author: sergey-miryanov

Lib/_pyrepl/reader.py

@@ -551,7 +551,10 @@ def get_prompt(self, lineno: int, cursor_on_line: bool) -> str:
         """Return what should be in the left-hand margin for line
         'lineno'."""
         if self.arg is not None and cursor_on_line:
-            prompt = self.__get_prompt_str(f"(arg: {self.arg}) ", DEFAULT_PS1)
+            prompt = DEFAULT_PS1
+            arg = self.__get_prompt_str(self.arg, "")
+            if arg:
+                prompt = f"(arg: {self.arg}) "
         elif self.paste_mode and not self.in_bracketed_paste:
             prompt = "(paste) "
         elif "\n" in self.buffer:

Lib/test/test_pyrepl/test_reader.py

@@ -281,15 +281,15 @@ def test_prompt_length(self):
         self.assertEqual(l, 5)
 
     def test_prompt_ps1_raise_exception(self):
-        # Handles simple ASCII prompt
+        # Handles exceptions from ps1 prompt
         class Prompt:
             def __str__(self): 1/0
 
         def prepare_reader_keep_prompts(*args, **kwargs):
             reader = prepare_reader(*args, **kwargs)
             del reader.get_prompt
             reader.ps1 = Prompt()
-            reader.ps2 = "+++ "
+            reader.ps2 = "... "
             reader.ps3 = "... "
             reader.ps4 = ""
             reader.can_colorize = False
@@ -305,23 +305,23 @@ def prepare_reader_keep_prompts(*args, **kwargs):
         prompt = reader.get_prompt(0, False)
         self.assertEqual(prompt, DEFAULT_PS1)
 
-    def test_prompt_ps2_raise_exception(self):
-        # Handles simple ASCII prompt
+    def test_prompt_ps2_ps3_ps4_raise_exception(self):
+        # Handles exceptions from ps2, ps3 and ps4 prompts
         class Prompt:
             def __str__(self): 1/0
 
         def prepare_reader_keep_prompts(*args, **kwargs):
             reader = prepare_reader(*args, **kwargs)
             del reader.get_prompt
-            reader.ps1 = "+++ "
+            reader.ps1 = Prompt()
             reader.ps2 = Prompt()
-            reader.ps3 = "--- "
-            reader.ps4 = "~~~ "
+            reader.ps3 = Prompt()
+            reader.ps4 = Prompt()
             reader.can_colorize = False
             reader.paste_mode = False
             return reader
 
-        events = code_to_events("if some_condition:\nsome_function()")
+        events = code_to_events("if some_condition:\nsome_function()\nsome_function()")
         reader, _ = handle_events_narrow_console(
             events,
             prepare_reader=prepare_reader_keep_prompts,
@@ -330,43 +330,22 @@ def prepare_reader_keep_prompts(*args, **kwargs):
         prompt = reader.get_prompt(0, False)
         self.assertEqual(prompt, DEFAULT_PS2)
 
-    def test_prompt_ps3_raise_exception(self):
-        # Handles simple ASCII prompt
-        class Prompt:
-            def __str__(self): 1/0
-
-        def prepare_reader_keep_prompts(*args, **kwargs):
-            reader = prepare_reader(*args, **kwargs)
-            del reader.get_prompt
-            reader.ps1 = "+++ "
-            reader.ps2 = "--- "
-            reader.ps3 = Prompt()
-            reader.ps4 = ""
-            reader.can_colorize = False
-            reader.paste_mode = False
-            return reader
-
-        events = code_to_events("if some_condition:\nsome_function()")
-        reader, _ = handle_events_narrow_console(
-            events,
-            prepare_reader=prepare_reader_keep_prompts,
-        )
-
         prompt = reader.get_prompt(1, False)
         self.assertEqual(prompt, DEFAULT_PS3)
 
-    def test_prompt_ps4_raise_exception(self):
-        # Handles simple ASCII prompt
+        prompt = reader.get_prompt(2, False)
+        self.assertEqual(prompt, DEFAULT_PS4)
+
+    def test_prompt_arg_raise_exception(self):
+        # Handles exceptions from arg prompt
         class Prompt:
             def __str__(self): 1/0
 
+            def __rmul__(self, b): return b
+
         def prepare_reader_keep_prompts(*args, **kwargs):
             reader = prepare_reader(*args, **kwargs)
             del reader.get_prompt
-            reader.ps1 = "+++ "
-            reader.ps2 = "--- "
-            reader.ps3 = "~~~ "
-            reader.ps4 = Prompt()
             reader.can_colorize = False
             reader.paste_mode = False
             return reader
@@ -377,8 +356,36 @@ def prepare_reader_keep_prompts(*args, **kwargs):
             prepare_reader=prepare_reader_keep_prompts,
         )
 
-        prompt = reader.get_prompt(1, False)
-        self.assertEqual(prompt, DEFAULT_PS4)
+        reader.arg = Prompt()
+        prompt = reader.get_prompt(0, True)
+        self.assertEqual(prompt, DEFAULT_PS1)
+
+    def test_prompt_raise_exception(self):
+        # Tests unrecoverable exceptions from prompts
+        cases = [
+            (MemoryError, "No memory for prompt"),
+            (SystemError, "System error for prompt"),
+        ]
+        for cls, msg in cases:
+            with self.subTest(msg):
+
+                class Prompt:
+                    def __str__(self): raise cls(msg)
+
+                def prepare_reader_keep_prompts(*args, **kwargs):
+                    reader = prepare_reader(*args, **kwargs)
+                    del reader.get_prompt
+                    reader.ps1 = Prompt()
+                    reader.can_colorize = False
+                    reader.paste_mode = False
+                    return reader
+
+                with self.assertRaisesRegex(cls, msg):
+                    events = code_to_events("a=1")
+                    handle_events_narrow_console(
+                        events,
+                        prepare_reader=prepare_reader_keep_prompts,
+                    )
 
     def test_completions_updated_on_key_press(self):
         namespace = {"itertools": itertools}