✨ feat: add apply_method_to_operator (#5)

Co-authored-by: Nyakku Shigure <sigure.qaq@gmail.com>

Author: zrr1999

src/expr_simplifier/transforms/auto_simplify.py

@@ -5,8 +5,11 @@
 from expr_simplifier.transforms.constant_folding import apply_constant_folding
 from expr_simplifier.transforms.cse import apply_cse
 from expr_simplifier.transforms.logical_simplification import apply_logical_simplification
+from expr_simplifier.transforms.method_to_operator import apply_method_to_operator
 from expr_simplifier.utils import loop_until_stable
 
 
 def auto_simplify(expr: ast.AST) -> ast.AST:
-    return loop_until_stable(expr, [apply_constant_folding, apply_logical_simplification, apply_cse], max_iter=100)
+    return loop_until_stable(
+        expr, [apply_constant_folding, apply_logical_simplification, apply_cse, apply_method_to_operator], max_iter=100
+    )

src/expr_simplifier/transforms/method_to_operator.py

@@ -0,0 +1,68 @@
+from __future__ import annotations
+
+import ast
+from ast import Add, BitAnd, BitOr, BitXor, Div, FloorDiv, Invert, LShift, Mod, Mult, Pow, RShift, Sub, UAdd, USub
+
+MAGIC_NAMES_TO_BINARY_OPS: dict[str, ast.operator] = {
+    "__add__": Add(),
+    "__sub__": Sub(),
+    "__mul__": Mult(),
+    "__truediv__": Div(),
+    "__floordiv__": FloorDiv(),
+    "__mod__": Mod(),
+    "__pow__": Pow(),
+    "__and__": BitAnd(),
+    "__or__": BitOr(),
+    "__xor__": BitXor(),
+    "__lshift__": LShift(),
+    "__rshift__": RShift(),
+}
+
+MAGIC_NAMES_TO_REVERSE_BINARY_OPS: dict[str, ast.operator] = {
+    "__radd__": Add(),
+    "__rsub__": Sub(),
+    "__rmul__": Mult(),
+    "__rtruediv__": Div(),
+    "__rfloordiv__": FloorDiv(),
+    "__rmod__": Mod(),
+    "__rpow__": Pow(),
+    "__rand__": BitAnd(),
+    "__ror__": BitOr(),
+    "__rxor__": BitXor(),
+    "__rlshift__": LShift(),
+    "__rrshift__": RShift(),
+}
+
+MAGIC_NAMES_TO_UNARY_OPS: dict[str, ast.unaryop] = {
+    "__neg__": USub(),
+    "__pos__": UAdd(),
+    "__invert__": Invert(),
+}
+
+
+class MethodToOperator(ast.NodeTransformer):
+    def visit_Call(self, node: ast.Call) -> ast.AST:
+        if not isinstance(node.func, ast.Attribute):
+            return self.generic_visit(node)
+        if node.func.attr in MAGIC_NAMES_TO_BINARY_OPS and len(node.args) == 1:
+            return ast.BinOp(
+                left=MethodToOperator().visit(node.func.value),
+                op=MAGIC_NAMES_TO_BINARY_OPS[node.func.attr],
+                right=MethodToOperator().visit(node.args[0]),
+            )
+        if node.func.attr in MAGIC_NAMES_TO_REVERSE_BINARY_OPS and len(node.args) == 1:
+            return ast.BinOp(
+                left=MethodToOperator().visit(node.args[0]),
+                op=MAGIC_NAMES_TO_REVERSE_BINARY_OPS[node.func.attr],
+                right=MethodToOperator().visit(node.func.value),
+            )
+        if node.func.attr in MAGIC_NAMES_TO_UNARY_OPS and len(node.args) == 0:
+            return ast.UnaryOp(
+                op=MAGIC_NAMES_TO_UNARY_OPS[node.func.attr],
+                operand=MethodToOperator().visit(node.func.value),
+            )
+        return self.generic_visit(node)
+
+
+def apply_method_to_operator(expr: ast.AST) -> ast.AST:
+    return MethodToOperator().visit(expr)

tests/test_transforms/test_method_to_operator.py

@@ -0,0 +1,72 @@
+from __future__ import annotations
+
+import ast
+
+import pytest
+
+from expr_simplifier.transforms.method_to_operator import apply_method_to_operator
+
+from .utils import check_expr_at_runtime
+
+
+@pytest.mark.parametrize(
+    ["expr", "expected"],
+    [
+        ("a.__add__(b)", "a + b"),
+        ("a.__sub__(b)", "a - b"),
+        ("a.__mul__(b)", "a * b"),
+        ("a.__truediv__(b)", "a / b"),
+        ("a.__floordiv__(b)", "a // b"),
+        ("a.__mod__(b)", "a % b"),
+        ("a.__pow__(b)", "a ** b"),
+        ("a.__and__(b)", "a & b"),
+        ("a.__or__(b)", "a | b"),
+        ("a.__xor__(b)", "a ^ b"),
+        ("a.__lshift__(b)", "a << b"),
+        ("a.__rshift__(b)", "a >> b"),
+        ("a.b.__add__(c.d)", "a.b + c.d"),
+        ("a.__add__(b).__mul__(c)", "(a + b) * c"),
+        ("a.__add__(b.__mul__(c))", "a + b * c"),
+        ("a.method(b).__add__(c)", "a.method(b) + c"),
+        ("a.__add__(b.method(c))", "a + b.method(c)"),
+        ("a.b.c.__add__(d.e.f)", "a.b.c + d.e.f"),
+        ("a.__add__(b).__sub__(c.__mul__(d))", "a + b - c * d"),
+    ],
+)
+def test_method_to_operator(expr: str, expected: str):
+    tree = ast.parse(expr, mode="eval")
+    transformed_tree = apply_method_to_operator(tree)
+    transformed_expr = ast.unparse(transformed_tree)
+    assert transformed_expr == expected
+    check_expr_at_runtime(tree, transformed_tree)
+
+
+@pytest.mark.parametrize(
+    ["expr", "expected"],
+    [
+        ("a.__neg__()", "-a"),
+        ("a.__pos__()", "+a"),
+        ("a.__invert__()", "~a"),
+    ],
+)
+def test_method_to_unary_operator(expr: str, expected: str):
+    tree = ast.parse(expr, mode="eval")
+    transformed_tree = apply_method_to_operator(tree)
+    transformed_expr = ast.unparse(transformed_tree)
+    assert transformed_expr == expected
+    check_expr_at_runtime(tree, transformed_tree)
+
+
+@pytest.mark.parametrize(
+    ["expr", "expected"],
+    [
+        ("a.method(b)", "a.method(b)"),
+        ("a.__unknown__(b)", "a.__unknown__(b)"),
+        ("a.__add__", "a.__add__"),
+    ],
+)
+def test_no_change(expr: str, expected: str):
+    tree = ast.parse(expr, mode="eval")
+    transformed_tree = apply_method_to_operator(tree)
+    transformed_expr = ast.unparse(transformed_tree)
+    assert transformed_expr == expected

tests/test_transforms/utils.py

@@ -20,6 +20,14 @@ def __radd__(self, other: Any):
         other = AnyObject.to_any_object(other)
         return AnyObject(f"{other.name} + {self.name}")
 
+    def __sub__(self, other: Any):
+        other = AnyObject.to_any_object(other)
+        return AnyObject(f"{self.name} - {other.name}")
+
+    def __rsub__(self, other: Any):
+        other = AnyObject.to_any_object(other)
+        return AnyObject(f"{other.name} - {self.name}")
+
     def __mul__(self, other: Any):
         other = AnyObject.to_any_object(other)
         return AnyObject(f"{self.name} * {other.name}")
@@ -44,6 +52,34 @@ def __rfloordiv__(self, other: Any):
         other = AnyObject.to_any_object(other)
         return AnyObject(f"{other.name} // {self.name}")
 
+    def __lshift__(self, other: Any):
+        other = AnyObject.to_any_object(other)
+        return AnyObject(f"{self.name} << {other.name}")
+
+    def __rshift__(self, other: Any):
+        other = AnyObject.to_any_object(other)
+        return AnyObject(f"{self.name} >> {other.name}")
+
+    def __or__(self, other: Any):
+        other = AnyObject.to_any_object(other)
+        return AnyObject(f"{self.name} | {other.name}")
+
+    def __xor__(self, other: Any):
+        other = AnyObject.to_any_object(other)
+        return AnyObject(f"{other.name} ^ {self.name}")
+
+    def __pow__(self, other: Any):
+        other = AnyObject.to_any_object(other)
+        return AnyObject(f"{self.name} ** {other.name}")
+
+    def __mod__(self, other: Any):
+        other = AnyObject.to_any_object(other)
+        return AnyObject(f"{self.name} % {other.name}")
+
+    def __and__(self, other: Any):
+        other = AnyObject.to_any_object(other)
+        return AnyObject(f"{self.name} & {other.name}")
+
     def __getattr__(self, name: str):
         return AnyObject(f"{self.name}.{name}")
 
@@ -81,6 +117,15 @@ def __le__(self, other: Any):
     def __repr__(self):
         return f"AnyObject({self.name})"
 
+    def __neg__(self):
+        return AnyObject(f"-{self.name}")
+
+    def __pos__(self):
+        return AnyObject(f"+{self.name}")
+
+    def __invert__(self):
+        return AnyObject(f"~{self.name}")
+
     @staticmethod
     def to_any_object(value: Any):
         if isinstance(value, AnyObject):