Improvement: Option to change collection concat to unpack

.gitignore

@@ -4,6 +4,7 @@
 .vscode/
 __pycache__/
 
+test.py
 *.ipynb
 
 build/

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
 

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
 

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),
+    )

version.txt

@@ -1 +1 @@
-7.0.0
+7.1.0