Merge pull request #994 from asottile/fstring-find-call

fix search for .encode(...) call with 3.12+ f-string tokens

pyupgrade/_plugins/default_encoding.py

@@ -12,12 +12,12 @@
 from pyupgrade._data import State
 from pyupgrade._data import TokenFunc
 from pyupgrade._string_helpers import is_codec
+from pyupgrade._token_helpers import find_call
 from pyupgrade._token_helpers import find_closing_bracket
-from pyupgrade._token_helpers import find_op
 
 
 def _fix_default_encoding(i: int, tokens: list[Token]) -> None:
-    i = find_op(tokens, i + 1, '(')
+    i = find_call(tokens, i + 1)
     j = find_closing_bracket(tokens, i)
     del tokens[i + 1:j]
 

pyupgrade/_token_helpers.py

@@ -48,6 +48,20 @@ def find_op(tokens: list[Token], i: int, src: str) -> int:
     return _find_token(tokens, i, 'OP', src)
 
 
+def find_call(tokens: list[Token], i: int) -> int:
+    depth = 0
+    while depth or not tokens[i].matches(name='OP', src='('):
+        if is_open(tokens[i]):  # pragma: >3.12 cover
+            depth += 1
+        elif is_close(tokens[i]):
+            # why max(...)? --
+            # ("something").method(...)
+            #  ^--start   target--^
+            depth = max(depth - 1, 0)
+        i += 1
+    return i
+
+
 def find_end(tokens: list[Token], i: int) -> int:
     while tokens[i].name != 'NEWLINE':
         i += 1

tests/features/default_encoding_test.py

@@ -38,6 +38,11 @@
             'f"{x}(".encode()',
             id='3.12+ handle open brace in fstring',
         ),
+        pytest.param(
+            'f"{foo(bar)}(".encode("utf-8")',
+            'f"{foo(bar)}(".encode()',
+            id='f-string with function call',
+        ),
     ),
 )
 def test_fix_encode(s, expected):