Skip to content

Commit 36f7005

Browse files
committed
fix search for .encode(...) call with 3.12+ f-string tokens
1 parent 7dc1fb3 commit 36f7005

File tree

3 files changed

+21
-2
lines changed

3 files changed

+21
-2
lines changed

pyupgrade/_plugins/default_encoding.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -12,12 +12,12 @@
1212
from pyupgrade._data import State
1313
from pyupgrade._data import TokenFunc
1414
from pyupgrade._string_helpers import is_codec
15+
from pyupgrade._token_helpers import find_call
1516
from pyupgrade._token_helpers import find_closing_bracket
16-
from pyupgrade._token_helpers import find_op
1717

1818

1919
def _fix_default_encoding(i: int, tokens: list[Token]) -> None:
20-
i = find_op(tokens, i + 1, '(')
20+
i = find_call(tokens, i + 1)
2121
j = find_closing_bracket(tokens, i)
2222
del tokens[i + 1:j]
2323

pyupgrade/_token_helpers.py

+14
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,20 @@ def find_op(tokens: list[Token], i: int, src: str) -> int:
4848
return _find_token(tokens, i, 'OP', src)
4949

5050

51+
def find_call(tokens: list[Token], i: int) -> int:
52+
depth = 0
53+
while depth or not tokens[i].matches(name='OP', src='('):
54+
if is_open(tokens[i]):
55+
depth += 1
56+
elif is_close(tokens[i]):
57+
# why max(...)? --
58+
# ("something").method(...)
59+
# ^--start target--^
60+
depth = max(depth - 1, 0)
61+
i += 1
62+
return i
63+
64+
5165
def find_end(tokens: list[Token], i: int) -> int:
5266
while tokens[i].name != 'NEWLINE':
5367
i += 1

tests/features/default_encoding_test.py

+5
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,11 @@
3838
'f"{x}(".encode()',
3939
id='3.12+ handle open brace in fstring',
4040
),
41+
pytest.param(
42+
'f"{foo(bar)}(".encode("utf-8")',
43+
'f"{foo(bar)}(".encode()',
44+
id='f-string with function call',
45+
),
4146
),
4247
)
4348
def test_fix_encode(s, expected):

0 commit comments

Comments
 (0)