From ddebae47e7a24e27882a5c48a099cc0dee6cb344 Mon Sep 17 00:00:00 2001 From: Alvin Wan Date: Sat, 4 Apr 2026 02:21:13 -0700 Subject: [PATCH 1/3] fix return simplifier crash --- pymini/pymini.py | 53 +++++++++++++++++++++++++++++------------------ tests/test_api.py | 33 +++++++++++++++++++++++++++++ 2 files changed, 66 insertions(+), 20 deletions(-) diff --git a/pymini/pymini.py b/pymini/pymini.py index ee5c66f..f5086ed 100644 --- a/pymini/pymini.py +++ b/pymini/pymini.py @@ -1,4 +1,5 @@ import ast +import copy import keyword from graphlib import TopologicalSorter from typing import Dict, List, Optional, Set @@ -41,39 +42,51 @@ class ReturnSimplifier(NodeTransformer): return (some code) - NOTE: unused_names must be modified in-place, since the set is passed to - RemoveUnusedVariables at initialization. Can't return a new set. + NOTE: unused_assignments must be modified in-place, since the set is passed + to RemoveUnusedVariables at initialization. Can't return a new set. """ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - self.name_to_node = {} - self.unused_names = set() + self.unused_assignments = set() - def visit_Assign(self, node: ast.Assign) -> ast.Assign: - if len(node.targets) == 1 and isinstance(node.targets[0], ast.Name): - self.name_to_node[node.targets[0].id] = node - return self.generic_visit(node) + def _can_simplify_return(self, previous: ast.stmt, current: ast.stmt) -> bool: + return ( + isinstance(previous, ast.Assign) + and len(previous.targets) == 1 + and isinstance(previous.targets[0], ast.Name) + and isinstance(current, ast.Return) + and isinstance(current.value, ast.Name) + and current.value.id == previous.targets[0].id + ) - def visit_Return(self, node: ast.Return) -> ast.Return: - if isinstance(node.value, ast.Name): - self.unused_names.add(node.value.id) - node = self.name_to_node[node.value.id] - return ast.Return(value=node.value) - return self.generic_visit(node) + def _simplify_body(self, body: List[ast.stmt]) -> List[ast.stmt]: + for previous, current in zip(body, body[1:]): + if self._can_simplify_return(previous, current): + self.unused_assignments.add(id(previous)) + current.value = copy.deepcopy(previous.value) + return body + + def generic_visit(self, node): + node = super().generic_visit(node) + for field, value in ast.iter_fields(node): + if isinstance(value, list) and value and all(isinstance(item, ast.stmt) for item in value): + setattr(node, field, self._simplify_body(value)) + return node class RemoveUnusedVariables(NodeTransformer): """Remove all unused variables. - NOTE: cannot store a copy of unused_names, as this set is modified in-place + NOTE: cannot store a copy of unused_assignments, as this set is modified + in-place after initialization. """ - def __init__(self, unused_names: Set[str]): + def __init__(self, unused_assignments: Set[int]): super().__init__() - self.unused_names = unused_names + self.unused_assignments = unused_assignments - def visit_Assign(self, node: ast.Name) -> ast.Name: - if isinstance(node.targets[0], ast.Name) and node.targets[0].id in self.unused_names: + def visit_Assign(self, node: ast.Assign) -> Optional[ast.Assign]: + if id(node) in self.unused_assignments: return None return self.generic_visit(node) @@ -812,7 +825,7 @@ def minify(sources, modules='main', keep_module_names=False, # simplify simplifier := ReturnSimplifier(), - RemoveUnusedVariables(simplifier.unused_names), + RemoveUnusedVariables(simplifier.unused_assignments), # minify ParentSetter(), diff --git a/tests/test_api.py b/tests/test_api.py index 2f4ab99..f67199d 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -80,6 +80,39 @@ def f(): assert modules == ["main"] +def test_minify_does_not_crash_when_returning_parameter_names(): + cleaned, modules = minify( + py( + """ + def abs_path(path): + if path: + return path + + value = 1 + return value + """ + ), + "main", + keep_global_variables=True, + keep_module_names=True, + ) + + tree = ast.parse(cleaned[0]) + function = next(node for node in tree.body if isinstance(node, ast.FunctionDef)) + condition = function.body[0] + simplified_return = function.body[1] + + assert isinstance(condition, ast.If) + assert isinstance(condition.body[0], ast.Return) + assert isinstance(condition.body[0].value, ast.Name) + assert condition.body[0].value.id == function.args.args[0].arg + + assert isinstance(simplified_return, ast.Return) + assert isinstance(simplified_return.value, ast.Constant) + assert simplified_return.value.value == 1 + assert modules == ["main"] + + def test_minify_updates_cross_file_imports(): cleaned, modules = minify( [ From 0f60bcd0dee1e24283fe4fc40dbf1a0dc8b281d0 Mon Sep 17 00:00:00 2001 From: Alvin Wan Date: Sat, 4 Apr 2026 02:32:07 -0700 Subject: [PATCH 2/3] add bare parameter return regression --- tests/test_api.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/tests/test_api.py b/tests/test_api.py index f67199d..0574405 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -113,6 +113,29 @@ def abs_path(path): assert modules == ["main"] +def test_minify_leaves_bare_parameter_returns_untouched(): + cleaned, modules = minify( + py( + """ + def f(path): + return path + """ + ), + "main", + keep_global_variables=True, + keep_module_names=True, + ) + + tree = ast.parse(cleaned[0]) + function = next(node for node in tree.body if isinstance(node, ast.FunctionDef)) + returned = function.body[0] + + assert isinstance(returned, ast.Return) + assert isinstance(returned.value, ast.Name) + assert returned.value.id == function.args.args[0].arg + assert modules == ["main"] + + def test_minify_updates_cross_file_imports(): cleaned, modules = minify( [ From 873a508ef3b1a44ed2da68ab9c2f7f420cdf050f Mon Sep 17 00:00:00 2001 From: Alvin Wan Date: Sat, 4 Apr 2026 02:39:30 -0700 Subject: [PATCH 3/3] remove redundant bare return test --- tests/test_api.py | 24 ------------------------ 1 file changed, 24 deletions(-) diff --git a/tests/test_api.py b/tests/test_api.py index 0574405..787e476 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -112,30 +112,6 @@ def abs_path(path): assert simplified_return.value.value == 1 assert modules == ["main"] - -def test_minify_leaves_bare_parameter_returns_untouched(): - cleaned, modules = minify( - py( - """ - def f(path): - return path - """ - ), - "main", - keep_global_variables=True, - keep_module_names=True, - ) - - tree = ast.parse(cleaned[0]) - function = next(node for node in tree.body if isinstance(node, ast.FunctionDef)) - returned = function.body[0] - - assert isinstance(returned, ast.Return) - assert isinstance(returned.value, ast.Name) - assert returned.value.id == function.args.args[0].arg - assert modules == ["main"] - - def test_minify_updates_cross_file_imports(): cleaned, modules = minify( [