Skip to content

Commit

Permalink
rewrite pep 646 Unpack to splat in *args
Browse files Browse the repository at this point in the history
  • Loading branch information
asottile committed Jul 31, 2023
1 parent b3e813e commit ccbbfcb
Show file tree
Hide file tree
Showing 2 changed files with 56 additions and 0 deletions.
39 changes: 39 additions & 0 deletions pyupgrade/_plugins/typing_pep646_unpack.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,3 +37,42 @@ def visit_Subscript(
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


def _visit_func(
state: State,
node: ast.AsyncFunctionDef | ast.FunctionDef,
parent: ast.AST,
) -> Iterable[tuple[Offset, TokenFunc]]:
if state.settings.min_version < (3, 11):
return

vararg = node.args.vararg
if (
vararg is not None and
isinstance(vararg.annotation, ast.Subscript) and
is_name_attr(
vararg.annotation.value,
state.from_imports,
('typing',), ('Unpack',),
)
):
yield ast_to_offset(vararg.annotation.value), _replace_unpack_with_star


@register(ast.AsyncFunctionDef)
def visit_AsyncFunctionDef(
state: State,
node: ast.AsyncFunctionDef,
parent: ast.AST,
) -> Iterable[tuple[Offset, TokenFunc]]:
yield from _visit_func(state, node, parent)


@register(ast.FunctionDef)
def visit_FunctionDef(
state: State,
node: ast.FunctionDef,
parent: ast.AST,
) -> Iterable[tuple[Offset, TokenFunc]]:
yield from _visit_func(state, node, parent)
17 changes: 17 additions & 0 deletions tests/features/typing_pep646_unpack_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,14 @@
' pass',
id='Not inside a subscript',
),
pytest.param(
'from typing import Unpack\n'
'from typing import TypedDict\n'
'class D(TypedDict):\n'
' x: int\n'
'def f(**kwargs: Unpack[D]) -> None: pass\n',
id='3.12 TypedDict for kwargs',
),
),
)
def test_fix_pep646_noop(s):
Expand Down Expand Up @@ -53,6 +61,15 @@ def test_fix_pep646_noop(s):
'class C(Generic[*Shape]):\n'
' pass',
),
pytest.param(
'from typing import Unpack\n'
'def f(*args: Unpack[tuple[int, ...]]): pass\n',
'from typing import Unpack\n'
'def f(*args: *tuple[int, ...]): pass\n',
id='Unpack for *args',
),
),
)
def test_typing_unpack(s, expected):
Expand Down

0 comments on commit ccbbfcb

Please sign in to comment.