diff --git a/.gitignore b/.gitignore index 5f28a6c..553074d 100644 --- a/.gitignore +++ b/.gitignore @@ -4,6 +4,7 @@ .vscode/ __pycache__/ +test.py *.ipynb build/ diff --git a/personal_python_ast_optimizer/parser/config.py b/personal_python_ast_optimizer/parser/config.py index 5ce59aa..6288550 100644 --- a/personal_python_ast_optimizer/parser/config.py +++ b/personal_python_ast_optimizer/parser/config.py @@ -139,6 +139,7 @@ class OptimizationsConfig(_Config): "functions_safe_to_exclude_in_test_expr", "remove_typing_cast", "remove_unused_imports", + "collection_concat_to_unpack", "fold_constants", "assume_this_machine", ) @@ -154,6 +155,7 @@ def __init__( # noqa: PLR0913 functions_safe_to_exclude_in_test_expr: set[str] | None = None, remove_unused_imports: bool = True, remove_typing_cast: bool = True, + collection_concat_to_unpack: bool = False, fold_constants: bool = False, assume_this_machine: bool = False, ) -> None: @@ -171,6 +173,7 @@ def __init__( # noqa: PLR0913 ) self.remove_unused_imports: bool = remove_unused_imports self.remove_typing_cast: bool = remove_typing_cast + self.collection_concat_to_unpack: bool = collection_concat_to_unpack self.assume_this_machine: bool = assume_this_machine self.fold_constants: bool = fold_constants diff --git a/personal_python_ast_optimizer/parser/skipper.py b/personal_python_ast_optimizer/parser/skipper.py index 4e64781..4f7538a 100644 --- a/personal_python_ast_optimizer/parser/skipper.py +++ b/personal_python_ast_optimizer/parser/skipper.py @@ -644,15 +644,40 @@ def visit_BinOp(self, node: ast.BinOp) -> ast.AST: # and constants need to be folded depth first parsed_node: ast.AST = self.generic_visit(node) - if ( - self.optimizations_config.fold_constants - and isinstance(parsed_node, ast.BinOp) - and isinstance(parsed_node.left, ast.Constant) - and isinstance(parsed_node.right, ast.Constant) - ): - return self._ast_constants_operation( - parsed_node.left, parsed_node.right, parsed_node.op - ) + if isinstance(parsed_node, ast.BinOp): + if ( + self.optimizations_config.fold_constants + and isinstance(parsed_node.left, ast.Constant) + and isinstance(parsed_node.right, ast.Constant) + ): + return self._ast_constants_operation( + parsed_node.left, parsed_node.right, parsed_node.op + ) + + if ( + self.optimizations_config.collection_concat_to_unpack + and isinstance(parsed_node.op, ast.Add) + and ( + isinstance(parsed_node.left, (ast.Tuple, ast.List)) + or isinstance(parsed_node.right, (ast.Tuple, ast.List)) + ) + ): + if ( + isinstance(parsed_node.left, ast.Tuple) + and isinstance(parsed_node.right, ast.Tuple) + or ( + isinstance(parsed_node.left, ast.List) + and isinstance(parsed_node.right, ast.List) + ) + ): # noqa: E721 + parsed_node.left.elts += parsed_node.right.elts + elif isinstance(parsed_node.left, (ast.Tuple, ast.List)): + parsed_node.left.elts.append(ast.Starred(parsed_node.right)) + else: + parsed_node.right.elts.insert(0, ast.Starred(parsed_node.left)) # type: ignore[attr-defined] + return parsed_node.right + + return parsed_node.left return parsed_node diff --git a/tests/parser/test_collections.py b/tests/parser/test_collections.py index 3eaa39e..de3683f 100644 --- a/tests/parser/test_collections.py +++ b/tests/parser/test_collections.py @@ -1,6 +1,9 @@ import pytest -from personal_python_ast_optimizer.parser.config import TokenTypesConfig +from personal_python_ast_optimizer.parser.config import ( + OptimizationsConfig, + TokenTypesConfig, +) from tests.utils import BeforeAndAfter, run_minifier_and_assert_correct @@ -92,3 +95,29 @@ class A(NamedTuple): before_and_after, token_types_config=TokenTypesConfig(simplify_named_tuples=True), ) + + +_collection_concat_to_unpack_test_cases: list[tuple[str, str]] = [ + ( + "a = (1,) + (0,0) + b", + "a=(1,0,0,*b)", + ), + ( + "a = [1] + b + [2]", + "a=[1,*b,2]", + ), + ( + "a = b + [2] + [3]", + "a=[*b,2,3]", + ), +] + + +@pytest.mark.parametrize(("before", "after"), _collection_concat_to_unpack_test_cases) +def test_collection_concat_to_unpack(before: str, after: str): + before_and_after = BeforeAndAfter(before, after) + + run_minifier_and_assert_correct( + before_and_after, + optimizations_config=OptimizationsConfig(collection_concat_to_unpack=True), + ) diff --git a/version.txt b/version.txt index 66ce77b..a3fcc71 100644 --- a/version.txt +++ b/version.txt @@ -1 +1 @@ -7.0.0 +7.1.0