_to_fstring: Use original token stream instead of unparsed AST

Signed-off-by: Anders Kaseorg <andersk@mit.edu>

pyupgrade/_main.py

@@ -35,8 +35,8 @@
 from pyupgrade._token_helpers import CLOSING
 from pyupgrade._token_helpers import KEYWORDS
 from pyupgrade._token_helpers import OPENING
+from pyupgrade._token_helpers import parse_call_args
 from pyupgrade._token_helpers import remove_brace
-from pyupgrade._token_helpers import victims
 
 DotFormatPart = Tuple[str, Optional[str], Optional[str], Optional[str]]
 
@@ -510,24 +510,12 @@ def _fix_tokens(contents_text: str, min_version: Version) -> str:
     return tokens_to_src(tokens).lstrip()
 
 
-def _simple_arg(arg: ast.expr) -> bool:
-    return (
-        isinstance(arg, ast.Name) or
-        (isinstance(arg, ast.Attribute) and _simple_arg(arg.value)) or
-        (
-            isinstance(arg, ast.Call) and
-            _simple_arg(arg.func) and
-            not arg.args and not arg.keywords
-        )
-    )
-
-
-def _format_params(call: ast.Call) -> Dict[str, str]:
-    params = {str(i): _unparse(arg) for i, arg in enumerate(call.args)}
+def _format_params(call: ast.Call) -> Set[str]:
+    params = {str(i) for i, arg in enumerate(call.args)}
     for kwd in call.keywords:
         # kwd.arg can't be None here because we exclude starargs
         assert kwd.arg is not None
-        params[kwd.arg] = _unparse(kwd.value)
+        params.add(kwd.arg)
     return params
 
 
@@ -566,8 +554,6 @@ def _parse(self, node: ast.Call) -> Optional[Tuple[DotFormatPart, ...]]:
                 isinstance(node.func, ast.Attribute) and
                 isinstance(node.func.value, ast.Str) and
                 node.func.attr == 'format' and
-                all(_simple_arg(arg) for arg in node.args) and
-                all(_simple_arg(k.value) for k in node.keywords) and
                 not has_starargs(node)
         ):
             return None
@@ -688,8 +674,6 @@ def _unparse(node: ast.expr) -> str:
         return node.id
     elif isinstance(node, ast.Attribute):
         return ''.join((_unparse(node.value), '.', node.attr))
-    elif isinstance(node, ast.Call):
-        return '{}()'.format(_unparse(node.func))
     elif isinstance(node, ast.Subscript):
         if sys.version_info >= (3, 9):  # pragma: no cover (py39+)
             node_slice: ast.expr = node.slice
@@ -704,7 +688,7 @@ def _unparse(node: ast.expr) -> str:
                 slice_s = ', '.join(_unparse(elt) for elt in node_slice.elts)
         else:
             slice_s = _unparse(node_slice)
-        return '{}[{}]'.format(_unparse(node.value), slice_s)
+        return f'{_unparse(node.value)}[{slice_s}]'
     elif isinstance(node, ast.Str):
         return repr(node.s)
     elif isinstance(node, ast.Ellipsis):
@@ -717,8 +701,28 @@ def _unparse(node: ast.expr) -> str:
         raise NotImplementedError(ast.dump(node))
 
 
-def _to_fstring(src: str, call: ast.Call) -> str:
-    params = _format_params(call)
+def _skip_unimportant_ws(tokens: List[Token], i: int) -> int:
+    while tokens[i].name == 'UNIMPORTANT_WS':
+        i += 1
+    return i
+
+
+def _to_fstring(
+    src: str, tokens: List[Token], args: List[Tuple[int, int]],
+) -> str:
+    params = {}
+    i = 0
+    for start, end in args:
+        start = _skip_unimportant_ws(tokens, start)
+        if tokens[start].name == 'NAME':
+            after = _skip_unimportant_ws(tokens, start + 1)
+            if tokens[after].src == '=':  # keyword argument
+                params[tokens[start].src] = tokens_to_src(
+                    tokens[after + 1:end],
+                ).strip()
+                continue
+        params[str(i)] = tokens_to_src(tokens[start:end]).strip()
+        i += 1
 
     parts = []
     i = 0
@@ -776,8 +780,6 @@ def _fix_py36_plus(contents_text: str) -> str:
         return contents_text
     for i, token in reversed_enumerate(tokens):
         if token.offset in visitor.fstrings:
-            node = visitor.fstrings[token.offset]
-
             # TODO: handle \N escape sequences
             if r'\N' in token.src:
                 continue
@@ -786,15 +788,19 @@ def _fix_py36_plus(contents_text: str) -> str:
             if tokens_to_src(tokens[i + 1:paren + 1]) != '.format(':
                 continue
 
-            # we don't actually care about arg position, so we pass `node`
-            fmt_victims = victims(tokens, paren, node, gen=False)
-            end = fmt_victims.ends[-1]
+            args, end = parse_call_args(tokens, paren)
             # if it spans more than one line, bail
-            if tokens[end].line != token.line:
+            if tokens[end - 1].line != token.line:
                 continue
 
-            tokens[i] = token._replace(src=_to_fstring(token.src, node))
-            del tokens[i + 1:end + 1]
+            args_src = tokens_to_src(tokens[paren:end])
+            if '\\' in args_src or '"' in args_src or "'" in args_src:
+                continue
+
+            tokens[i] = token._replace(
+                src=_to_fstring(token.src, tokens, args),
+            )
+            del tokens[i + 1:end]
         elif token.offset in visitor.named_tuples and token.name == 'NAME':
             call = visitor.named_tuples[token.offset]
             types: Dict[str, ast.expr] = {

pyupgrade/_plugins/oserror_aliases.py

@@ -62,7 +62,7 @@ def _fix_oserror_except(
     if len(unique_args) > 1:
         joined = '({})'.format(', '.join(unique_args))
     elif tokens[start - 1].name != 'UNIMPORTANT_WS':
-        joined = ' {}'.format(unique_args[0])
+        joined = f' {unique_args[0]}'
     else:
         joined = unique_args[0]
 

tests/features/fstrings_test.py

@@ -30,6 +30,10 @@
         r'"\N{snowman} {}".format(a)',
         # not enough placeholders / placeholders missing
         '"{}{}".format(a)', '"{a}{b}".format(a=a)',
+        # backslashes and quotes cannot nest
+        r'''"{}".format(a['\\'])''',
+        '"{}".format(a["b"])',
+        "'{}'.format(a['b'])",
     ),
 )
 def test_fix_fstrings_noop(s):
@@ -50,6 +54,7 @@ def test_fix_fstrings_noop(s):
         ('"hello {}!".format(name)', 'f"hello {name}!"'),
         ('"{}{{}}{}".format(escaped, y)', 'f"{escaped}{{}}{y}"'),
         ('"{}{b}{}".format(a, c, b=b)', 'f"{a}{b}{c}"'),
+        ('"{}".format(0x0)', 'f"{0x0}"'),
         # TODO: poor man's f-strings?
         # '"{foo}".format(**locals())'
     ),