Add simple constant folding

dflook · dflook · commit 5b9ae1b5b7f3 · 2023-08-16T08:45:44.000+01:00
diff --git a/src/python_minifier/__init__.py b/src/python_minifier/__init__.py
@@ -20,6 +20,7 @@
 )
 
 from python_minifier.transforms.combine_imports import CombineImports
+from python_minifier.transforms.constant_folding import FoldConstants
 from python_minifier.transforms.remove_annotations import RemoveAnnotations
 from python_minifier.transforms.remove_annotations_options import RemoveAnnotationsOptions
 from python_minifier.transforms.remove_asserts import RemoveAsserts
@@ -69,7 +70,8 @@ def minify(
     remove_asserts=False,
     remove_debug=False,
     remove_explicit_return_none=True,
-    remove_builtin_exception_brackets=True
+    remove_builtin_exception_brackets=True,
+    constant_folding=True
 ):
     """
     Minify a python module
@@ -102,6 +104,7 @@ def minify(
     :param bool remove_debug: If conditional statements that test '__debug__ is True' should be removed
     :param bool remove_explicit_return_none: If explicit return None statements should be replaced with a bare return
     :param bool remove_builtin_exception_brackets: If brackets should be removed when raising exceptions with no arguments
+    :param bool constant_folding: If constant expressions should be evaluated
 
     :rtype: str
 
@@ -150,6 +153,9 @@ def minify(
     if remove_explicit_return_none:
         module = RemoveExplicitReturnNone()(module)
 
+    if constant_folding:
+        module = FoldConstants()(module)
+
     bind_names(module)
     resolve_names(module)
 
diff --git a/src/python_minifier/transforms/constant_folding.py b/src/python_minifier/transforms/constant_folding.py
@@ -0,0 +1,55 @@
+import ast
+
+from python_minifier.expression_printer import ExpressionPrinter
+from python_minifier.transforms.suite_transformer import SuiteTransformer
+from python_minifier.util import is_ast_node
+
+
+class FoldConstants(SuiteTransformer):
+    """
+    Fold Constants if it would reduce the size of the source
+    """
+
+    def __init__(self):
+        super(FoldConstants, self).__init__()
+
+    def visit_BinOp(self, node):
+
+        node.left = self.visit(node.left)
+        node.right = self.visit(node.right)
+
+        # Check this is a constant expression that could be folded
+        if not is_ast_node(node.left, (ast.Num, ast.Str, ast.Bytes, ast.NameConstant)):
+            return node
+        if not is_ast_node(node.right, (ast.Num, ast.Str, ast.Bytes, ast.NameConstant)):
+            return node
+
+        expression_printer = ExpressionPrinter()
+
+        try:
+            original_expression = expression_printer(node)
+            value = eval(original_expression)
+        except Exception as e:
+            return node
+
+        if isinstance(value, str):
+            new_node = ast.Str(s=value)
+        elif isinstance(value, bytes):
+            new_node = ast.Bytes(s=value)
+        elif isinstance(value, bool):
+            new_node = ast.NameConstant(value=value)
+        elif isinstance(value, (int, float, complex)):
+            new_node = ast.Num(n=value)
+        else:
+            return node
+
+        expression_printer = ExpressionPrinter()
+        folded_expression = expression_printer(new_node)
+
+        if len(folded_expression) > len(original_expression):
+            # Result is longer than original expression
+            return node
+
+        assert eval(folded_expression) == value
+
+        return self.add_child(new_node, node.parent, node.namespace)
