From 73c323c9522cc42757d1d9a2e8a6f7272d551e0a Mon Sep 17 00:00:00 2001 From: Anders Kaseorg Date: Mon, 25 May 2020 22:07:24 -0700 Subject: [PATCH] _to_fstring: Use original token stream instead of unparsed AST Signed-off-by: Anders Kaseorg --- pyupgrade/_main.py | 68 +++++++++++++++------------ pyupgrade/_plugins/oserror_aliases.py | 2 +- tests/features/fstrings_test.py | 5 ++ 3 files changed, 43 insertions(+), 32 deletions(-) diff --git a/pyupgrade/_main.py b/pyupgrade/_main.py index e3035001..ca1542fb 100644 --- a/pyupgrade/_main.py +++ b/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] = { diff --git a/pyupgrade/_plugins/oserror_aliases.py b/pyupgrade/_plugins/oserror_aliases.py index e0d5a704..d7cb9bbd 100644 --- a/pyupgrade/_plugins/oserror_aliases.py +++ b/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] diff --git a/tests/features/fstrings_test.py b/tests/features/fstrings_test.py index 89d32a7d..791c3a22 100644 --- a/tests/features/fstrings_test.py +++ b/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())' ),