diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index bcbd274a..244df129 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -34,7 +34,7 @@ repos: - id: add-trailing-comma args: [--py36-plus] - repo: https://github.com/asottile/pyupgrade - rev: v2.18.1 + rev: v2.19.0 hooks: - id: pyupgrade args: [--py36-plus] diff --git a/README.md b/README.md index 6004942b..8786487b 100644 --- a/README.md +++ b/README.md @@ -20,7 +20,7 @@ Sample `.pre-commit-config.yaml`: ```yaml - repo: https://github.com/asottile/pyupgrade - rev: v2.18.1 + rev: v2.19.0 hooks: - id: pyupgrade ``` @@ -440,6 +440,17 @@ Availability: ``` +### Unpacking argument list comprehensions + +Availability: +- `--py3-plus` is passed on the commandline. + +```diff +-foo(*[i for i in bar]) ++foo(*(i for i in bar)) +``` + + ### `typing.NamedTuple` / `typing.TypedDict` py36+ syntax Availability: diff --git a/pyupgrade/_plugins/generator_expressions_pep289.py b/pyupgrade/_plugins/generator_expressions_pep289.py index e9169d61..e042d8a3 100644 --- a/pyupgrade/_plugins/generator_expressions_pep289.py +++ b/pyupgrade/_plugins/generator_expressions_pep289.py @@ -3,6 +3,7 @@ from typing import List from typing import Tuple +from tokenize_rt import NON_CODING_TOKENS from tokenize_rt import Offset from tokenize_rt import Token @@ -12,11 +13,10 @@ from pyupgrade._data import TokenFunc from pyupgrade._token_helpers import find_closing_bracket from pyupgrade._token_helpers import find_comprehension_opening_bracket +from pyupgrade._token_helpers import replace_list_comp_brackets ALLOWED_FUNCS = frozenset(( - 'all', - 'any', 'bytearray', 'bytes', 'frozenset', @@ -34,13 +34,11 @@ def _delete_list_comp_brackets(i: int, tokens: List[Token]) -> None: end = find_closing_bracket(tokens, start) tokens[end] = Token('PLACEHOLDER', '') tokens[start] = Token('PLACEHOLDER', '') - - -def _replace_list_comp_brackets(i: int, tokens: List[Token]) -> None: - start = find_comprehension_opening_bracket(i, tokens) - end = find_closing_bracket(tokens, start) - tokens[end] = Token('OP', ')') - tokens[start] = Token('OP', '(') + j = end + 1 + while j < len(tokens) and tokens[j].name in NON_CODING_TOKENS: + j += 1 + if tokens[j].name == 'OP' and tokens[j].src == ',': + tokens[j] = Token('PLACEHOLDER', '') def _func_condition(func: ast.expr) -> bool: @@ -71,4 +69,4 @@ def visit_Call( if len(node.args) == 1 and not node.keywords: yield ast_to_offset(node.args[0]), _delete_list_comp_brackets else: - yield ast_to_offset(node.args[0]), _replace_list_comp_brackets + yield ast_to_offset(node.args[0]), replace_list_comp_brackets diff --git a/pyupgrade/_plugins/unpacking_argument_list_comprehensions.py b/pyupgrade/_plugins/unpacking_argument_list_comprehensions.py new file mode 100644 index 00000000..7b94032f --- /dev/null +++ b/pyupgrade/_plugins/unpacking_argument_list_comprehensions.py @@ -0,0 +1,24 @@ +import ast +from typing import Iterable +from typing import Tuple + +from tokenize_rt import Offset + +from pyupgrade._ast_helpers import ast_to_offset +from pyupgrade._data import register +from pyupgrade._data import State +from pyupgrade._data import TokenFunc +from pyupgrade._token_helpers import replace_list_comp_brackets + + +@register(ast.Starred) +def visit_Starred( + state: State, + node: ast.Starred, + parent: ast.AST, +) -> Iterable[Tuple[Offset, TokenFunc]]: + if ( + state.settings.min_version >= (3,) and + isinstance(node.value, ast.ListComp) + ): + yield ast_to_offset(node.value), replace_list_comp_brackets diff --git a/pyupgrade/_token_helpers.py b/pyupgrade/_token_helpers.py index c37117dd..e031972a 100644 --- a/pyupgrade/_token_helpers.py +++ b/pyupgrade/_token_helpers.py @@ -474,3 +474,10 @@ def find_comprehension_opening_bracket(i: int, tokens: List[Token]) -> int: return i else: # pragma: no cover ( None: + start = find_comprehension_opening_bracket(i, tokens) + end = find_closing_bracket(tokens, start) + tokens[end] = Token('OP', ')') + tokens[start] = Token('OP', '(') diff --git a/setup.cfg b/setup.cfg index f0880fb0..9e442a27 100644 --- a/setup.cfg +++ b/setup.cfg @@ -1,6 +1,6 @@ [metadata] name = pyupgrade -version = 2.18.1 +version = 2.19.0 description = A tool to automatically upgrade syntax for newer versions. long_description = file: README.md long_description_content_type = text/markdown diff --git a/tests/features/generator_expressions_pep289_test.py b/tests/features/generator_expressions_pep289_test.py index ae032823..1709ec7e 100644 --- a/tests/features/generator_expressions_pep289_test.py +++ b/tests/features/generator_expressions_pep289_test.py @@ -78,6 +78,28 @@ def test_fix_typing_text_noop(s): id='Join function', ), + pytest.param( + '"".join([[i for _ in range(2)] for i in range(3)],)\n', + + '"".join([i for _ in range(2)] for i in range(3))\n', + + id='Trailing comma after list comprehension', + ), + pytest.param( + 'sum(\n' + ' [\n' + ' i for i in range(3)\n' + ' ],\n' + ')\n', + + 'sum(\n' + ' \n' + ' i for i in range(3)\n' + ' \n' + ')\n', + + id='Multiline list comprehension with trailing comma\n', + ), ), ) def test_fix_typing_text(s, expected): diff --git a/tests/features/unpacking_argument_list_comprehensions_test.py b/tests/features/unpacking_argument_list_comprehensions_test.py new file mode 100644 index 00000000..ce4487c9 --- /dev/null +++ b/tests/features/unpacking_argument_list_comprehensions_test.py @@ -0,0 +1,75 @@ +import pytest + +from pyupgrade._data import Settings +from pyupgrade._main import _fix_plugins + + +@pytest.mark.parametrize( + ('s', 'version'), + ( + pytest.param( + 'foo(*[i for i in bar])\n', + (2, 7), + id='Not Python3+', + ), + pytest.param( + '2*3', + (3,), + id='Multiplication star', + ), + pytest.param( + '2**3', + (3,), + id='Power star', + ), + pytest.param( + 'foo([i for i in bar])', + (3,), + id='List comp, no star', + ), + pytest.param( + 'foo(*bar)', + (3,), + id='Starred, no list comp', + ), + ), +) +def test_fix_unpack_argument_list_comp_noop(s, version): + assert _fix_plugins(s, settings=Settings(min_version=version)) == s + + +@pytest.mark.parametrize( + ('s', 'expected'), + ( + pytest.param( + 'foo(*[i for i in bar])\n', + + 'foo(*(i for i in bar))\n', + + id='Starred list comprehension', + ), + pytest.param( + 'foo(\n' + ' *\n' + ' [i for i in bar]\n' + ' )\n', + + 'foo(\n' + ' *\n' + ' (i for i in bar)\n' + ' )\n', + + id='Multiline starred list comprehension', + ), + pytest.param( + 'foo(*[i for i in bar], qux, quox=None)\n', + + 'foo(*(i for i in bar), qux, quox=None)\n', + + id='Single line, including other args', + ), + ), +) +def test_fix_unpack_argument_list_comp(s, expected): + ret = _fix_plugins(s, settings=Settings((3,))) + assert ret == expected