diff --git a/pyupgrade/_plugins/typing_pep646_unpack.py b/pyupgrade/_plugins/typing_pep646_unpack.py new file mode 100644 index 00000000..1259c60b --- /dev/null +++ b/pyupgrade/_plugins/typing_pep646_unpack.py @@ -0,0 +1,39 @@ +from __future__ import annotations + +import ast +from typing import Iterable + +from tokenize_rt import Offset +from tokenize_rt import Token + +from pyupgrade._ast_helpers import ast_to_offset +from pyupgrade._ast_helpers import is_name_attr +from pyupgrade._data import register +from pyupgrade._data import State +from pyupgrade._data import TokenFunc +from pyupgrade._token_helpers import find_closing_bracket +from pyupgrade._token_helpers import find_token +from pyupgrade._token_helpers import remove_brace + + +def _replace_unpack_with_star(i: int, tokens: list[Token]) -> None: + start = find_token(tokens, i, '[') + end = find_closing_bracket(tokens, start) + + remove_brace(tokens, end) + # replace `Unpack` with `*` + tokens[i:start + 1] = [tokens[i]._replace(name='OP', src='*')] + + +@register(ast.Subscript) +def visit_Subscript( + state: State, + node: ast.Subscript, + parent: ast.AST, +) -> Iterable[tuple[Offset, TokenFunc]]: + if state.settings.min_version < (3, 11): + return + + if is_name_attr(node.value, state.from_imports, ('typing',), ('Unpack',)): + if isinstance(parent, (ast.Subscript, ast.Index)): + yield ast_to_offset(node.value), _replace_unpack_with_star diff --git a/tests/features/typing_pep646_unpack_test.py b/tests/features/typing_pep646_unpack_test.py new file mode 100644 index 00000000..c1d82e85 --- /dev/null +++ b/tests/features/typing_pep646_unpack_test.py @@ -0,0 +1,60 @@ +from __future__ import annotations + +import pytest + +from pyupgrade._data import Settings +from pyupgrade._main import _fix_plugins + + +@pytest.mark.parametrize( + ('s',), + ( + pytest.param( + 'from typing import Unpack\n' + 'foo(Unpack())', + id='Not a subscript', + ), + pytest.param( + 'from typing import TypeVarTuple, Unpack\n' + 'Shape = TypeVarTuple("Shape")\n' + 'class Foo(Unpack[Shape]):\n' + ' pass', + id='Not inside a subscript', + ), + ), +) +def test_fix_pep646_noop(s): + assert _fix_plugins(s, settings=Settings(min_version=(3, 11))) == s + assert _fix_plugins(s, settings=Settings(min_version=(3, 10))) == s + + +@pytest.mark.parametrize( + ('s', 'expected'), + ( + ( + 'from typing import Generic, TypeVarTuple, Unpack\n' + "Shape = TypeVarTuple('Shape')\n" + 'class C(Generic[Unpack[Shape]]):\n' + ' pass', + + 'from typing import Generic, TypeVarTuple, Unpack\n' + "Shape = TypeVarTuple('Shape')\n" + 'class C(Generic[*Shape]):\n' + ' pass', + ), + ( + 'from typing import Generic, TypeVarTuple, Unpack\n' + "Shape = TypeVarTuple('Shape')\n" + 'class C(Generic[Unpack [Shape]]):\n' + ' pass', + + 'from typing import Generic, TypeVarTuple, Unpack\n' + "Shape = TypeVarTuple('Shape')\n" + 'class C(Generic[*Shape]):\n' + ' pass', + ), + ), +) +def test_typing_unpack(s, expected): + assert _fix_plugins(s, settings=Settings(min_version=(3, 11))) == expected + assert _fix_plugins(s, settings=Settings(min_version=(3, 10))) == s