More fixes to f-string parsing

openrewrite · Aug 27, 2024 · 27c5696 · 27c5696
1 parent 112cd24
commit 27c5696
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 25 deletions.
diff --git a/rewrite/poetry.lock b/rewrite/poetry.lock
diff --git a/rewrite/pyproject.toml b/rewrite/pyproject.toml
@@ -11,6 +11,7 @@ packages = [
 [tool.poetry.dependencies]
 python = ">=3.9"
 rewrite-remote = { path = "../../../moderneinc/rewrite-remote/python/rewrite-remote", develop = true }
+more-itertools = "^10.4.0"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^8.3.2"

diff --git a/rewrite/rewrite/python/_parser_visitor.py b/rewrite/rewrite/python/_parser_visitor.py
@@ -1,5 +1,6 @@
 import ast
 import token
+from more_itertools import peekable
 from functools import lru_cache
 from io import BytesIO
 from pathlib import Path
@@ -897,7 +898,7 @@ def visit_JoinedStr(self, node):
         tokens = tokenize(BytesIO(self._source[self._cursor:].encode('utf-8')).readline)
         next(tokens)  # skip ENCODING token
         tok = next(tokens)  # FSTRING_START token
-        return self.__map_fstring(node, prefix, tok, tokens)[0]
+        return self.__map_fstring(node, prefix, tok, peekable(tokens))[0]
 
     def visit_FormattedValue(self, node):
         raise ValueError("This method should not be called directly")
@@ -987,7 +988,7 @@ def visit_Module(self, node: ast.Module) -> py.CompilationUnit:
                 self.__pad_right(j.Empty(random_id(), Space.EMPTY, Markers.EMPTY), Space.EMPTY)],
             self.__whitespace()
         )
-        # assert self._cursor == len(self._source)
+        assert self._cursor == len(self._source)
         return cu
 
     def visit_Name(self, node):
@@ -1351,7 +1352,7 @@ def _map_assignment_operator(self, op):
             raise ValueError(f"Unsupported operator: {op}")
         return self.__pad_left(self.__source_before(op_str), op)
 
-    def __map_fstring(self, node: ast.JoinedStr, prefix: Space, tok: TokenInfo, tokens):
+    def __map_fstring(self, node: ast.JoinedStr, prefix: Space, tok: TokenInfo, tokens: peekable):
         if tok.type != token.FSTRING_START:
             if len(node.values) == 1 and isinstance(node.values[0], ast.Constant):
                 # format specifiers are stored as f-strings in the AST; e.g. `f'{1:n}'`
@@ -1387,41 +1388,33 @@ def __map_fstring(self, node: ast.JoinedStr, prefix: Space, tok: TokenInfo, toke
                         nested,
                         Space.EMPTY
                     )
-                    prev_tok = tok
-                    self._cursor += len(prev_tok.string)
-                    tok = next(tokens)
                 else:
                     expr = self.__pad_right(
                         self.__convert(cast(ast.FormattedValue, value).value),
                         self.__whitespace()
                     )
-                    prev_tok = tok
                     try:
-                        while (tok := next(tokens)).type not in (token.FSTRING_END, token.FSTRING_MIDDLE):
-                            prev_tok = tok
-                            if prev_tok.type == token.OP and prev_tok.string == '!':
+                        while (tokens.peek()).type not in (token.FSTRING_END, token.FSTRING_MIDDLE):
+                            tok = next(tokens)
+                            if tok.type == token.OP and tok.string == '!':
                                 break
                     except StopIteration:
                         pass
-                    self._cursor += len(prev_tok.string)
 
                 # conversion specifier
-                if prev_tok.type == token.OP and prev_tok.string == '!':
+                if tok.type == token.OP and tok.string == '!':
+                    self._cursor += len(tok.string)
                     tok = next(tokens)
                     conv = py.FormattedString.Value.Conversion.ASCII if tok.string == 'a' else py.FormattedString.Value.Conversion.STR if tok.string == 's' else py.FormattedString.Value.Conversion.REPR
                     self._cursor += len(tok.string)
-                    prev_tok = next(tokens)
                     tok = next(tokens)
-                    self._cursor += len(tok.string)
                 else:
                     conv = None
 
                 # format specifier
-                if prev_tok.type == token.OP and prev_tok.string == ':':
-                    format_spec, tok = self.__map_fstring(cast(ast.JoinedStr, cast(ast.FormattedValue, value).format_spec), Space.EMPTY, tok, tokens)
-                    # self._cursor += len(tok.string)
-                    # tok = next(tokens)
-                    # self._cursor += len(tok.string)
+                if tok.type == token.OP and tok.string == ':':
+                    self._cursor += len(tok.string)
+                    format_spec, tok = self.__map_fstring(cast(ast.JoinedStr, cast(ast.FormattedValue, value).format_spec), Space.EMPTY, next(tokens), tokens)
                 else:
                     format_spec = None
                 parts.append(py.FormattedString.Value(
@@ -1432,6 +1425,8 @@ def __map_fstring(self, node: ast.JoinedStr, prefix: Space, tok: TokenInfo, toke
                     conv,
                     format_spec
                 ))
+                self._cursor += len(tok.string)
+                tok = next(tokens)
             else:  # FSTRING_MIDDLE
                 save_cursor = self._cursor
                 while True:
@@ -1451,6 +1446,8 @@ def __map_fstring(self, node: ast.JoinedStr, prefix: Space, tok: TokenInfo, toke
         if consume_end_delim:
             self._cursor += len(tok.string)  # FSTRING_END token
             tok = next(tokens)
+        elif tok.type == token.FSTRING_MIDDLE and len(tok.string) == 0:
+            tok = next(tokens)
 
         return (py.FormattedString(
             random_id(),

diff --git a/rewrite/tests/python/all/fstring_test.py b/rewrite/tests/python/all/fstring_test.py
@@ -76,7 +76,9 @@ def test_debug():
 def test_conversion():
     # language=python
     rewrite_run(python("""a = f'{"foo"!a}'"""))
+    # language=python
     rewrite_run(python("""a = f'{"foo"!s}'"""))
+    # language=python
     rewrite_run(python("""a = f'{"foo"!r}'"""))
 
 
@@ -90,7 +92,6 @@ def test_conversion_and_format_expr():
     rewrite_run(python("""a = f'{"foo"!s:<{5*2}}'"""))
 
 
-@pytest.mark.xfail(reason="Implementation still not quite correct", strict=True)
 def test_nested_fstring_conversion_and_format_expr():
     # language=python
     rewrite_run(python("""a = f'{f"foo"!s:<{5*2}}'"""))
@@ -130,11 +131,9 @@ def test_nested_fstring_format():
 
 def test_format_value():
     # language=python
-    # rewrite_run(python("a = f'{1:.{2 + 3}f}'"))
+    rewrite_run(python("a = f'{1:.{2 + 3}f}'"))
     # language=python
-    # rewrite_run(python('''a = f"{'abc':>{2*3}}"'''))
-    # language=python
-    rewrite_run(python("""a = f'{f"{'foo'}":>{2*3}}'"""))
+    rewrite_run(python('''a = f"{'abc':>{2*3}}"'''))
 
 
 def test_nested_fstring_with_format_value():