From f1e94ab838df3f6cf84ccc72fad6bfbc4ae0d99c Mon Sep 17 00:00:00 2001
From: Delgan <delgan.py@gmail.com>
Date: Sun, 10 Sep 2023 15:19:46 +0200
Subject: [PATCH] Fix f-string formatting in traceback of Python 3.12

In Python 3.12, new tokens were added to "tokenize" module in order to
differentiate simple strings and f-string. Additionally, the expressions
inside the f-string are properly parsed as well. The unit tests have
been updated consequently.
---
 loguru/_better_exceptions.py                | 20 ++++++++++++++------
 tests/exceptions/output/modern/f_string.txt | 13 ++++++++++---
 tests/exceptions/source/modern/f_string.py  |  6 +++++-
 tests/test_exceptions_formatting.py         |  2 +-
 4 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/loguru/_better_exceptions.py b/loguru/_better_exceptions.py
index de44bd54..8327a13b 100644
--- a/loguru/_better_exceptions.py
+++ b/loguru/_better_exceptions.py
@@ -46,6 +46,12 @@ class SyntaxHighlighter:
     _builtins = set(dir(builtins))
     _constants = {"True", "False", "None"}
     _punctation = {"(", ")", "[", "]", "{", "}", ":", ",", ";"}
+    _strings = {tokenize.STRING}
+    _fstring_middle = None
+
+    if sys.version_info >= (3, 12):
+        _strings.update({tokenize.FSTRING_START, tokenize.FSTRING_MIDDLE, tokenize.FSTRING_END})
+        _fstring_middle = tokenize.FSTRING_MIDDLE
 
     def __init__(self, style=None):
         self._style = style or self._default_style
@@ -56,7 +62,12 @@ def highlight(self, source):
         output = ""
 
         for token in self.tokenize(source):
-            type_, string, start, end, line = token
+            type_, string, (start_row, start_column), (_, end_column), line = token
+
+            if type_ == self._fstring_middle:
+                # When an f-string contains "{{" or "}}", they appear as "{" or "}" in the "string"
+                # attribute of the token. However, they do not count in the column position.
+                end_column += string.count("{") + string.count("}")
 
             if type_ == tokenize.NAME:
                 if string in self._constants:
@@ -74,23 +85,20 @@ def highlight(self, source):
                     color = style["operator"]
             elif type_ == tokenize.NUMBER:
                 color = style["number"]
-            elif type_ == tokenize.STRING:
+            elif type_ in self._strings:
                 color = style["string"]
             elif type_ == tokenize.COMMENT:
                 color = style["comment"]
             else:
                 color = style["other"]
 
-            start_row, start_column = start
-            _, end_column = end
-
             if start_row != row:
                 source = source[column:]
                 row, column = start_row, 0
 
             if type_ != tokenize.ENCODING:
                 output += line[column:start_column]
-                output += color.format(string)
+                output += color.format(line[start_column:end_column])
 
             column = end_column
 
diff --git a/tests/exceptions/output/modern/f_string.txt b/tests/exceptions/output/modern/f_string.txt
index a657d254..97415415 100644
--- a/tests/exceptions/output/modern/f_string.txt
+++ b/tests/exceptions/output/modern/f_string.txt
@@ -1,11 +1,18 @@
 
 Traceback (most recent call last):
 
-  File "tests/exceptions/source/modern/f_string.py", line 17, in <module>
+  File "tests/exceptions/source/modern/f_string.py", line 21, in <module>
     hello()
     └ <function hello at 0xDEADBEEF>
 
-  File "tests/exceptions/source/modern/f_string.py", line 13, in hello
-    f"{name}" and f'{{ {f / 0} }}'
+  File "tests/exceptions/source/modern/f_string.py", line 11, in hello
+    output = f"Hello" + f' ' + f"""World""" and world()
+                                                └ <function world at 0xDEADBEEF>
+
+  File "tests/exceptions/source/modern/f_string.py", line 17, in world
+    f"{name} -> { f }" and {} or f'{{ {f / 0} }}'
+       │          │                    └ 1
+       │          └ 1
+       └ 'world'
 
 ZeroDivisionError: division by zero
diff --git a/tests/exceptions/source/modern/f_string.py b/tests/exceptions/source/modern/f_string.py
index 43b67bc9..339a41ae 100644
--- a/tests/exceptions/source/modern/f_string.py
+++ b/tests/exceptions/source/modern/f_string.py
@@ -8,9 +8,13 @@
 
 
 def hello():
+    output = f"Hello" + f' ' + f"""World""" and world()
+
+
+def world():
     name = "world"
     f = 1
-    f"{name}" and f'{{ {f / 0} }}'
+    f"{name} -> { f }" and {} or f'{{ {f / 0} }}'
 
 
 with logger.catch():
diff --git a/tests/test_exceptions_formatting.py b/tests/test_exceptions_formatting.py
index b2f97b21..a24375cc 100644
--- a/tests/test_exceptions_formatting.py
+++ b/tests/test_exceptions_formatting.py
@@ -230,7 +230,6 @@ def test_exception_others(filename):
     "filename, minimum_python_version",
     [
         ("type_hints", (3, 6)),
-        ("f_string", (3, 6)),
         ("positional_only_argument", (3, 8)),
         ("walrus_operator", (3, 8)),
         ("match_statement", (3, 10)),
@@ -242,6 +241,7 @@ def test_exception_others(filename):
         ("grouped_as_cause_and_context", (3, 11)),
         ("grouped_max_length", (3, 11)),
         ("grouped_max_depth", (3, 11)),
+        ("f_string", (3, 12)),  # Available since 3.6 but in 3.12 the lexer for f-string changed.
     ],
 )
 def test_exception_modern(filename, minimum_python_version):