From fe42aaca7b000200ea9915c04c01f2fb3f4dbaf6 Mon Sep 17 00:00:00 2001
From: Georg Brandl <georg@python.org>
Date: Fri, 3 Feb 2023 07:47:18 +0100
Subject: [PATCH] Python console: make traceback handling more robust

Fixes #2329
Fixes #2226
---
 CHANGES                            |  1 +
 pygments/lexers/python.py          | 24 ++++----
 tests/snippets/pycon/broken_tb.txt | 97 ++++++++++++++++++++++++++++++
 3 files changed, 110 insertions(+), 12 deletions(-)
 create mode 100644 tests/snippets/pycon/broken_tb.txt

diff --git a/CHANGES b/CHANGES
index e101317840..445dace88b 100644
--- a/CHANGES
+++ b/CHANGES
@@ -21,6 +21,7 @@ Version 2.15.0
   on Python 3.10 and older (#2328).
 - Fix some places where a locale-dependent encoding could unintentionally
   be used instead of UTF-8 (#2326).
+- Fix Python traceback handling (#2226, #2329).
 
 Version 2.14.0
 --------------
diff --git a/pygments/lexers/python.py b/pygments/lexers/python.py
index 0a318a9ed4..6ba3d19ec6 100644
--- a/pygments/lexers/python.py
+++ b/pygments/lexers/python.py
@@ -679,15 +679,15 @@ def get_tokens_unprocessed(self, text):
         insertions = []
         curtb = ''
         tbindex = 0
-        tb = 0
+        in_tb = False
         for match in line_re.finditer(text):
             line = match.group()
             if line.startswith('>>> ') or line.startswith('... '):
-                tb = 0
+                in_tb = False
                 insertions.append((len(curcode),
                                    [(0, Generic.Prompt, line[:4])]))
                 curcode += line[4:]
-            elif line.rstrip() == '...' and not tb:
+            elif line.rstrip() == '...' and not in_tb:
                 # only a new >>> prompt can end an exception block
                 # otherwise an ellipsis in place of the traceback frames
                 # will be mishandled
@@ -700,20 +700,20 @@ def get_tokens_unprocessed(self, text):
                         insertions, pylexer.get_tokens_unprocessed(curcode))
                     curcode = ''
                     insertions = []
-                if (line.startswith('Traceback (most recent call last):') or
-                        re.match('  File "[^"]+", line \\d+\\n$', line)):
-                    tb = 1
-                    curtb = line
-                    tbindex = match.start()
-                elif line == 'KeyboardInterrupt\n':
-                    yield match.start(), Name.Class, line
-                elif tb:
+                if in_tb:
                     curtb += line
                     if not (line.startswith(' ') or line.strip() == '...'):
-                        tb = 0
+                        in_tb = False
                         for i, t, v in tblexer.get_tokens_unprocessed(curtb):
                             yield tbindex+i, t, v
                         curtb = ''
+                elif (line.startswith('Traceback (most recent call last):') or
+                        re.match('  File "[^"]+", line \\d+\\n$', line)):
+                    in_tb = True
+                    curtb = line
+                    tbindex = match.start()
+                elif line == 'KeyboardInterrupt\n':
+                    yield match.start(), Name.Class, line
                 else:
                     yield match.start(), Generic.Output, line
         if curcode:
diff --git a/tests/snippets/pycon/broken_tb.txt b/tests/snippets/pycon/broken_tb.txt
new file mode 100644
index 0000000000..6cecf7612b
--- /dev/null
+++ b/tests/snippets/pycon/broken_tb.txt
@@ -0,0 +1,97 @@
+---input---
+>>> exec('"')
+Traceback (most recent call last):
+  File "<stdin>", line 1, in <module>
+  File "<string>", line 1
+    "
+     ^
+SyntaxError: EOL while scanning string literal
+
+>>> exec('"')
+Traceback (most recent call last):
+  File "<stdin>", line 1, in <module>
+  File "<string>", line 1, in <module>
+    "
+     ^
+SyntaxError: EOL while scanning string literal
+
+---tokens---
+'>>> '        Generic.Prompt
+'exec'        Name
+'('           Punctuation
+"'"           Literal.String.Single
+'"'           Literal.String.Single
+"'"           Literal.String.Single
+')'           Punctuation
+'\n'          Text.Whitespace
+
+'Traceback (most recent call last):\n' Generic.Traceback
+
+'  File '     Text
+'"<stdin>"'   Name.Builtin
+', line '     Text
+'1'           Literal.Number
+', in '       Text
+'<module>'    Name
+'\n'          Text.Whitespace
+
+'  File '     Text
+'"<string>"'  Name.Builtin
+', line '     Text
+'1'           Literal.Number
+'\n'          Text.Whitespace
+
+'    '        Text.Whitespace
+'"'           Literal.String.Double
+'\n'          Text.Whitespace
+
+'     '       Text.Whitespace
+'^'           Punctuation.Marker
+'\n'          Text.Whitespace
+
+'SyntaxError' Generic.Error
+': '          Text
+'EOL while scanning string literal' Name
+'\n'          Text.Whitespace
+
+'\n'          Generic.Output
+
+'>>> '        Generic.Prompt
+'exec'        Name
+'('           Punctuation
+"'"           Literal.String.Single
+'"'           Literal.String.Single
+"'"           Literal.String.Single
+')'           Punctuation
+'\n'          Text.Whitespace
+
+'Traceback (most recent call last):\n' Generic.Traceback
+
+'  File '     Text
+'"<stdin>"'   Name.Builtin
+', line '     Text
+'1'           Literal.Number
+', in '       Text
+'<module>'    Name
+'\n'          Text.Whitespace
+
+'  File '     Text
+'"<string>"'  Name.Builtin
+', line '     Text
+'1'           Literal.Number
+', in '       Text
+'<module>'    Name
+'\n'          Text.Whitespace
+
+'    '        Text.Whitespace
+'"'           Literal.String.Double
+'\n'          Text.Whitespace
+
+'     '       Text.Whitespace
+'^'           Punctuation.Marker
+'\n'          Text.Whitespace
+
+'SyntaxError' Generic.Error
+': '          Text
+'EOL while scanning string literal' Name
+'\n'          Text.Whitespace