Also support folding complex with negative real

python · Mar 5, 2023 · 3bae7d9 · 3bae7d9
1 parent 3b3c4c7
commit 3bae7d9
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 5 deletions.
diff --git a/mypy/constant_fold.py b/mypy/constant_fold.py
@@ -88,9 +88,9 @@ def constant_fold_binary_op(
 
     if isinstance(left, float) and isinstance(right, float):
         return constant_fold_binary_float_op(op, left, right)
-    if isinstance(left, float) and isinstance(right, int):
+    elif isinstance(left, float) and isinstance(right, int):
         return constant_fold_binary_float_op(op, left, right)
-    if isinstance(left, int) and isinstance(right, float):
+    elif isinstance(left, int) and isinstance(right, float):
         return constant_fold_binary_float_op(op, left, right)
 
     if op == "+" and isinstance(left, str) and isinstance(right, str):
@@ -104,6 +104,10 @@ def constant_fold_binary_op(
         return left + right
     elif op == "+" and isinstance(left, complex) and isinstance(right, (int, float)):
         return left + right
+    elif op == "-" and isinstance(left, (int, float)) and isinstance(right, complex):
+        return left - right
+    elif op == "-" and isinstance(left, complex) and isinstance(right, (int, float)):
+        return left - right
 
     return None
 
@@ -148,7 +152,7 @@ def constant_fold_binary_float_op(op: str, left: int | float, right: int | float
     assert not (isinstance(left, int) and isinstance(right, int)), (op, left, right)
     if op == "+":
         return left + right
-    if op == "-":
+    elif op == "-":
         return left - right
     elif op == "*":
         return left * right

diff --git a/mypyc/test-data/fixtures/ir.py b/mypyc/test-data/fixtures/ir.py
@@ -130,6 +130,7 @@ def __init__(self, x: object, y: object = None) -> None: pass
     def __add__(self, n: complex) -> complex: pass
     def __radd__(self, n: float) -> complex: pass
     def __sub__(self, n: complex) -> complex: pass
+    def __rsub__(self, n: float) -> complex: pass
     def __mul__(self, n: complex) -> complex: pass
     def __truediv__(self, n: complex) -> complex: pass
     def __neg__(self) -> complex: pass

diff --git a/mypyc/test-data/irbuild-constant-fold.test b/mypyc/test-data/irbuild-constant-fold.test
@@ -431,24 +431,36 @@ FLOAT_N: Final = 1.5
 def integral() -> None:
     pos = 1+2j
     pos_2 = 2j+N
+    neg = 1-2j
+    neg_2 = 2j-N
 def floating() -> None:
     pos = 1.5+2j
     pos_2 = 2j+FLOAT_N
+    neg = 1.5-2j
+    neg_2 = 2j-FLOAT_N
 [out]
 def integral():
-    r0, pos, r1, pos_2 :: object
+    r0, pos, r1, pos_2, r2, neg, r3, neg_2 :: object
 L0:
     r0 = (1+2j)
     pos = r0
     r1 = (1+2j)
     pos_2 = r1
+    r2 = (1-2j)
+    neg = r2
+    r3 = (-1+2j)
+    neg_2 = r3
     return 1
 def floating():
-    r0, pos, r1, pos_2 :: object
+    r0, pos, r1, pos_2, r2, neg, r3, neg_2 :: object
 L0:
     r0 = (1.5+2j)
     pos = r0
     r1 = (1.5+2j)
     pos_2 = r1
+    r2 = (1.5-2j)
+    neg = r2
+    r3 = (-1.5+2j)
+    neg_2 = r3
     return 1