Correctly fold UnaryOp nodes

Author: dflook

src/python_minifier/transforms/constant_folding.py

@@ -10,36 +10,35 @@
 from python_minifier.util import is_constant_node
 
 
-class FoldConstants(SuiteTransformer):
-    """
-    Fold Constants if it would reduce the size of the source
+def is_foldable_constant(node):
     """
+    Check if a node is a constant expression that can participate in folding.
 
-    def __init__(self):
-        super(FoldConstants, self).__init__()
+    We can asume that children have already been folded, so foldable constants are either:
+    - Simple literals (Num, NameConstant)
+    - UnaryOp(USub/Invert) on a Num - these don't fold to shorter forms,
+      so they remain after child visiting. UAdd and Not would have been
+      folded away since they always produce shorter results.
+    """
+    if is_constant_node(node, (ast.Num, ast.NameConstant)):
+        return True
 
-    def visit_BinOp(self, node):
+    if isinstance(node, ast.UnaryOp):
+        if isinstance(node.op, (ast.USub, ast.Invert)):
+            return is_constant_node(node.operand, ast.Num)
 
-        node.left = self.visit(node.left)
-        node.right = self.visit(node.right)
+    return False
 
-        # Check this is a constant expression that could be folded
-        # We don't try to fold strings or bytes, since they have probably been arranged this way to make the source shorter and we are unlikely to beat that
-        if not is_constant_node(node.left, (ast.Num, ast.NameConstant)):
-            return node
-        if not is_constant_node(node.right, (ast.Num, ast.NameConstant)):
-            return node
 
-        if isinstance(node.op, ast.Div):
-            # Folding div is subtle, since it can have different results in Python 2 and Python 3
-            # Do this once target version options have been implemented
-            return node
+class FoldConstants(SuiteTransformer):
+    """
+    Fold Constants if it would reduce the size of the source
+    """
 
-        if isinstance(node.op, ast.Pow):
-            # This can be folded, but it is unlikely to reduce the size of the source
-            # It can also be slow to evaluate
-            return node
+    def __init__(self):
+        super(FoldConstants, self).__init__()
 
+    def fold(self, node):
         # Evaluate the expression
         try:
             original_expression = unparse_expression(node)
@@ -96,6 +95,44 @@ def visit_BinOp(self, node):
         # New representation is shorter and has the same value, so use it
         return self.add_child(new_node, get_parent(node), node.namespace)
 
+    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
+        # We don't try to fold strings or bytes, since they have probably been arranged this way to make the source shorter and we are unlikely to beat that
+        if not is_foldable_constant(node.left):
+            return node
+        if not is_foldable_constant(node.right):
+            return node
+
+        if isinstance(node.op, ast.Div):
+            # Folding div is subtle, since it can have different results in Python 2 and Python 3
+            # Do this once target version options have been implemented
+            return node
+
+        if isinstance(node.op, ast.Pow):
+            # This can be folded, but it is unlikely to reduce the size of the source
+            # It can also be slow to evaluate
+            return node
+
+        return self.fold(node)
+
+    def visit_UnaryOp(self, node):
+
+        node.operand = self.visit(node.operand)
+
+        # Only fold if the operand is a foldable constant
+        if not is_foldable_constant(node.operand):
+            return node
+
+        # Only fold these unary operators
+        if not isinstance(node.op, (ast.USub, ast.UAdd, ast.Invert, ast.Not)):
+            return node
+
+        return self.fold(node)
+
 
 def equal_value_and_type(a, b):
     if type(a) != type(b):

test/test_folding.py

@@ -3,11 +3,11 @@
 
 import pytest
 
+from python_minifier import minify
 from python_minifier.ast_annotation import add_parent
 from python_minifier.ast_compare import compare_ast
 from python_minifier.rename import add_namespace
-from python_minifier.transforms.constant_folding import FoldConstants
-
+from python_minifier.transforms.constant_folding import FoldConstants, equal_value_and_type
 
 def fold_constants(source):
     module = ast.parse(source)
@@ -106,7 +106,9 @@ def test_bool(source, expected):
         ('0xf0|0x0f', '0xff'),
         ('10%2', '0'),
         ('10%3', '1'),
-        ('10-100', '-90')
+        ('10-100', '-90'),
+        ('1+1', '2'),
+        ('2+2', '4'),
     ]
 )
 def test_int(source, expected):
@@ -149,3 +151,262 @@ def test_not_eval(source, expected):
     """
 
     run_test(source, expected)
+
+
+@pytest.mark.parametrize(
+    ('source', 'expected'), [
+        # Folding results in infinity, which can be represented as 1e999
+        ('1e308 + 1e308', '1e999'),
+        ('1e308 * 2', '1e999'),
+    ]
+)
+def test_fold_infinity(source, expected):
+    """
+    Test that expressions resulting in infinity are folded to 1e999.
+
+    Infinity can be represented as 1e999, which is shorter than
+    the original expression.
+    """
+    run_test(source, expected)
+
+
+@pytest.mark.parametrize(
+    ('source', 'expected'), [
+        # Folding would result in NaN, which cannot be represented as a literal
+        ('1e999 - 1e999', '1e999 - 1e999'),
+        ('0.0 * 1e999', '0.0 * 1e999'),
+    ]
+)
+def test_no_fold_nan(source, expected):
+    """
+    Test that expressions resulting in NaN are not folded.
+
+    NaN is not a valid Python literal, so we cannot fold expressions
+    that would produce it.
+    """
+    run_test(source, expected)
+
+
+@pytest.mark.parametrize(
+    ('source', 'expected'), [
+        ('100.0+100.0', '200.0'),
+        ('1000.0+1000.0', '2000.0'),
+    ]
+)
+def test_fold_float(source, expected):
+    """
+    Test that float expressions are folded when the result is shorter.
+    """
+    run_test(source, expected)
+
+
+def test_equal_value_and_type():
+    """
+    Test the equal_value_and_type helper function.
+    """
+
+    # Same type and value
+    assert equal_value_and_type(1, 1) is True
+    assert equal_value_and_type(1.0, 1.0) is True
+    assert equal_value_and_type(True, True) is True
+    assert equal_value_and_type('hello', 'hello') is True
+
+    # Different types
+    assert equal_value_and_type(1, 1.0) is False
+    assert equal_value_and_type(1, True) is False
+    assert equal_value_and_type(True, 1) is False
+
+    # Different values
+    assert equal_value_and_type(1, 2) is False
+    assert equal_value_and_type(1.0, 2.0) is False
+
+
+def test_equal_value_and_type_nan():
+    """
+    Test the equal_value_and_type helper function with NaN values.
+    """
+
+    nan = float('nan')
+
+    # NaN is not equal to itself in Python (nan != nan is True)
+    # But if both are NaN, equal_value_and_type returns True via a == b
+    # Since nan == nan is False, we need to check the actual behavior
+    result = equal_value_and_type(nan, nan)
+    # Python's nan == nan is False, so this should be False
+    assert result is False
+
+    # NaN compared to non-NaN should be False
+    assert equal_value_and_type(nan, 1.0) is False
+    assert equal_value_and_type(1.0, nan) is False
+
+
+@pytest.mark.parametrize(
+    ('source', 'expected'), [
+        ('5 - 10', '-5'),
+        ('0 - 100', '-100'),
+        ('1.0 - 2.0', '-1.0'),
+        ('0.0 - 100.0', '-100.0'),
+    ]
+)
+def test_negative_results(source, expected):
+    """
+    Test BinOp expressions that produce negative results.
+    """
+    run_test(source, expected)
+
+
+@pytest.mark.parametrize(
+    ('source', 'expected'), [
+        ('5 * -2', '-10'),
+        ('-5 * 2', '-10'),
+        ('-5 + 10', '5'),
+        ('-90 + 10', '-80'),
+        ('10 - 20 + 5', '-5'),
+        ('(5 - 10) * 2', '-10'),
+        ('2 * (0 - 5)', '-10'),
+        ('(1 - 10) + (2 - 20)', '-27'),
+    ]
+)
+def test_negative_operands_folded(source, expected):
+    """
+    Test that expressions with negative operands are folded.
+    """
+    run_test(source, expected)
+
+
+@pytest.mark.parametrize(
+    ('source', 'expected'), [
+        ('-(-5)', '5'),
+        ('--5', '5'),
+        ('-(-100)', '100'),
+        ('-(-(5 + 5))', '10'),
+        ('~(~0)', '0'),
+        ('~~5', '5'),
+        ('~~100', '100'),
+        ('+(+5)', '5'),
+        ('+(-5)', '-5'),
+    ]
+)
+def test_unary_folded(source, expected):
+    """
+    Test that unary operations on constant expressions are folded.
+    """
+    run_test(source, expected)
+
+
+@pytest.mark.parametrize(
+    ('source', 'expected'), [
+        ('not not True', 'True'),
+        ('not not False', 'False'),
+        ('not True', 'False'),
+        ('not False', 'True'),
+    ]
+)
+def test_unary_not_folded(source, expected):
+    """
+    Test that 'not' operations on constant expressions are folded.
+    """
+    if sys.version_info < (3, 4):
+        pytest.skip('NameConstant not in python < 3.4')
+
+    run_test(source, expected)
+
+
+@pytest.mark.parametrize(
+    ('source', 'expected'), [
+        ('-5', '-5'),
+        ('~5', '~5'),
+    ]
+)
+def test_unary_simple_not_folded(source, expected):
+    """
+    Test that simple unary operations on literals are not folded
+    when the result would not be shorter.
+    """
+    run_test(source, expected)
+
+
+def test_unary_plus_folded():
+    """
+    Test that unary plus on a literal is folded to remove the plus.
+    """
+    run_test('+5', '5')
+
+
+def test_constant_folding_enabled_by_default():
+    """Verify constant folding is enabled by default."""
+    source = 'x = 10 + 10'
+    result = minify(source)
+    assert '20' in result
+    assert '10+10' not in result and '10 + 10' not in result
+
+
+def test_constant_folding_disabled():
+    """Verify expressions are not folded when constant_folding=False."""
+    source = 'x = 10 + 10'
+    result = minify(source, constant_folding=False)
+    assert '10+10' in result or '10 + 10' in result
+    assert result.strip() != 'x=20'
+
+
+def test_constant_folding_disabled_complex_expression():
+    """Verify complex expressions are preserved when disabled."""
+    source = 'SECONDS_IN_A_DAY = 60 * 60 * 24'
+    result = minify(source, constant_folding=False)
+    assert '60*60*24' in result or '60 * 60 * 24' in result
+
+
+def test_constant_folding_enabled_complex_expression():
+    """Verify complex expressions are folded when enabled."""
+    source = 'SECONDS_IN_A_DAY = 60 * 60 * 24'
+    result = minify(source, constant_folding=True)
+    assert '86400' in result
+
+
+@pytest.mark.parametrize(
+    ('source', 'should_contain_when_disabled'), [
+        ('x = 5 - 10', '5-10'),
+        ('x = True | False', 'True|False'),
+        ('x = 0xff ^ 0x0f', '255^15'),
+    ]
+)
+def test_constant_folding_disabled_various_ops(source, should_contain_when_disabled):
+    """Verify various operations are not folded when disabled."""
+    if sys.version_info < (3, 4) and 'True' in source:
+        pytest.skip('NameConstant not in python < 3.4')
+
+    result = minify(source, constant_folding=False)
+    assert should_contain_when_disabled in result.replace(' ', '')
+
+
+@pytest.mark.parametrize(
+    ('source', 'expected'), [
+        ('1j + 2j', '3j'),
+        ('3j * 2', '6j'),
+        ('2 * 3j', '6j'),
+        ('10j - 5j', '5j'),
+    ]
+)
+def test_complex_folded(source, expected):
+    """
+    Test complex number operations that are folded.
+
+    Complex operations are folded when the result is shorter than the original.
+    """
+    run_test(source, expected)
+
+
+@pytest.mark.parametrize(
+    ('source', 'expected'), [
+        ('1j - 2j', '1j-2j'),
+        ('1j * 1j', '1j*1j'),
+        ('0j + 5', '0j+5'),
+        ('2j / 1j', '2j/1j'),
+        ('1j ** 2', '1j**2'),
+    ]
+)
+def test_complex_not_folded(source, expected):
+    """
+    Test complex number operations that are not folded.
+    """
+    run_test(source, expected)