-
-
Notifications
You must be signed in to change notification settings - Fork 170
/
fstrings.py
139 lines (116 loc) · 4.23 KB
/
fstrings.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
from __future__ import annotations
import ast
from typing import Iterable
from tokenize_rt import Offset
from tokenize_rt import parse_string_literal
from tokenize_rt import Token
from tokenize_rt import tokens_to_src
from pyupgrade._ast_helpers import ast_to_offset
from pyupgrade._ast_helpers import contains_await
from pyupgrade._ast_helpers import has_starargs
from pyupgrade._data import register
from pyupgrade._data import State
from pyupgrade._data import TokenFunc
from pyupgrade._string_helpers import parse_format
from pyupgrade._string_helpers import unparse_parsed_string
from pyupgrade._token_helpers import parse_call_args
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
# need to remove `u` prefix so it isn't invalid syntax
prefix, rest = parse_string_literal(src)
new_src = 'f' + prefix.translate({ord('u'): None, ord('U'): None}) + rest
for s, name, spec, conv in parse_format(new_src):
if name is not None:
k, dot, rest = name.partition('.')
name = ''.join((params[k or str(i)], dot, rest))
if not k: # named and auto params can be in different orders
i += 1
parts.append((s, name, spec, conv))
return unparse_parsed_string(parts)
def _fix_fstring(i: int, tokens: list[Token]) -> None:
token = tokens[i]
paren = i + 3
if tokens_to_src(tokens[i + 1:paren + 1]) != '.format(':
return
args, end = parse_call_args(tokens, paren)
# if it spans more than one line, bail
if tokens[end - 1].line != token.line:
return
args_src = tokens_to_src(tokens[paren:end])
if '\\' in args_src or '"' in args_src or "'" in args_src:
return
tokens[i] = token._replace(src=_to_fstring(token.src, tokens, args))
del tokens[i + 1:end]
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.add(kwd.arg)
return params
@register(ast.Call)
def visit_Call(
state: State,
node: ast.Call,
parent: ast.AST,
) -> Iterable[tuple[Offset, TokenFunc]]:
if state.settings.min_version < (3, 6):
return
if (
isinstance(node.func, ast.Attribute) and
isinstance(node.func.value, ast.Str) and
node.func.attr == 'format' and
not has_starargs(node)
):
try:
parsed = parse_format(node.func.value.s)
except ValueError:
return
params = _format_params(node)
seen = set()
i = 0
for _, name, spec, _ in parsed:
# timid: difficult to rewrite correctly
if spec is not None and '{' in spec:
break
if name is not None:
candidate, _, _ = name.partition('.')
# timid: could make the f-string longer
if candidate and candidate in seen:
break
# timid: bracketed
elif '[' in name:
break
seen.add(candidate)
key = candidate or str(i)
# their .format() call is broken currently
if key not in params:
break
if not candidate:
i += 1
else:
if (
state.settings.min_version >= (3, 7) or
not contains_await(node)
):
yield ast_to_offset(node), _fix_fstring