Skip to content

Commit

Permalink
Add PEP646 Unpack plugin
Browse files Browse the repository at this point in the history
Co-authored-by: Marco Edward Gorelli <33491632+MarcoGorelli@users.noreply.github.com>
  • Loading branch information
2 people authored and asottile committed Oct 29, 2022
1 parent b1f3615 commit 848751f
Show file tree
Hide file tree
Showing 2 changed files with 99 additions and 0 deletions.
39 changes: 39 additions & 0 deletions 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
60 changes: 60 additions & 0 deletions 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

0 comments on commit 848751f

Please sign in to comment.