Merge pull request #745 from tusharsadhwani/unpack

asottile · web-flow · commit 1268511b2d0e · 2022-10-29T17:50:38.000-04:00
Add PEP646 Unpack plugin
diff --git a/pyupgrade/_plugins/typing_pep646_unpack.py 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
@@ -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
