[flows] adding support for math expressions

pywa/types/flows.py

@@ -1325,7 +1325,7 @@ def __str__(self) -> str:
 
     @staticmethod
     def _format_value(val: Ref | bool | int | float | str) -> str:
-        if isinstance(val, Ref):
+        if isinstance(val, (Ref, MathExpression)):
             return val.to_str()
         elif isinstance(val, str):
             return f"'{val}'"
@@ -1351,6 +1351,36 @@ def __lt__(self, other: Ref | int | float) -> Condition:
     def __le__(self, other: Ref | int | float) -> Condition:
         return Condition(f"({self.to_str()} <= {self._format_value(other)})")
 
+    def __add__(self, other: Ref | int | float) -> MathExpression:
+        return MathExpression(f"({self.to_str()} + {self._format_value(other)})")
+
+    def __radd__(self, other: Ref | int | float) -> MathExpression:
+        return MathExpression(f"({self._format_value(other)} + {self.to_str()})")
+
+    def __sub__(self, other: Ref | int | float) -> MathExpression:
+        return MathExpression(f"({self.to_str()} - {self._format_value(other)})")
+
+    def __rsub__(self, other: Ref | int | float) -> MathExpression:
+        return MathExpression(f"({self._format_value(other)} - {self.to_str()})")
+
+    def __mul__(self, other: Ref | int | float) -> MathExpression:
+        return MathExpression(f"({self.to_str()} * {self._format_value(other)})")
+
+    def __rmul__(self, other: Ref | int | float) -> MathExpression:
+        return MathExpression(f"({self._format_value(other)} * {self.to_str()})")
+
+    def __truediv__(self, other: Ref | int | float) -> MathExpression:
+        return MathExpression(f"({self.to_str()} / {self._format_value(other)})")
+
+    def __rtruediv__(self, other: Ref | int | float) -> MathExpression:
+        return MathExpression(f"({self._format_value(other)} / {self.to_str()})")
+
+    def __mod__(self, other: Ref | int | float) -> MathExpression:
+        return MathExpression(f"({self.to_str()} % {self._format_value(other)})")
+
+    def __rmod__(self, other: Ref | int | float) -> MathExpression:
+        return MathExpression(f"({self._format_value(other)} % {self.to_str()})")
+
     def __and__(self, other: Ref | Condition) -> Condition:
         if isinstance(other, Ref):
             return Condition(f"({self.to_str()} && {other.to_str()})")
@@ -1371,6 +1401,20 @@ def __repr__(self) -> str:
 _RefT = TypeVar("_RefT", bound=Ref)
 
 
+class MathExpression:
+    def __init__(self, expression: str):
+        self._expression = expression
+
+    def __str__(self) -> str:
+        return self.to_str()
+
+    def __repr__(self) -> str:
+        return f"MathExpression({self._expression})"
+
+    def to_str(self) -> str:
+        return self._expression
+
+
 class Condition:
     def __init__(self, expression: str):
         self._expression = expression
@@ -1388,12 +1432,15 @@ def __or__(self, other: Condition | Ref) -> Condition:
     def __invert__(self) -> Condition:
         return Condition(f"!{self._expression}")
 
-    def to_str(self) -> str:
-        return self._expression
+    def __str__(self) -> str:
+        return self.to_str()
 
     def __repr__(self) -> str:
         return f"Condition({self._expression})"
 
+    def to_str(self) -> str:
+        return self._expression
+
 
 class ScreenDataRef(Ref):
     """

tests/test_flows.py

@@ -157,6 +157,97 @@ def test_ref_less_than_or_equal():
     assert condition.to_str() == "(${data.age} <= 21)"
 
 
+def test_ref_addition():
+    ref = Ref(prefix="data", field="age")
+    math_op = ref + 21
+    assert math_op.to_str() == "(${data.age} + 21)"
+
+
+def test_ref_right_addition():
+    ref = Ref(prefix="data", field="age")
+    math_op = 21 + ref
+    assert math_op.to_str() == "(21 + ${data.age})"
+
+
+def test_ref_subtraction():
+    ref = Ref(prefix="data", field="age")
+    math_op = ref - 21
+    assert math_op.to_str() == "(${data.age} - 21)"
+
+
+def test_ref_right_subtraction():
+    ref = Ref(prefix="data", field="age")
+    math_op = 21 - ref
+    assert math_op.to_str() == "(21 - ${data.age})"
+
+
+def test_ref_multiplication():
+    ref = Ref(prefix="data", field="age")
+    math_op = ref * 21
+    assert math_op.to_str() == "(${data.age} * 21)"
+
+
+def test_ref_right_multiplication():
+    ref = Ref(prefix="data", field="age")
+    math_op = 21 * ref
+    assert math_op.to_str() == "(21 * ${data.age})"
+
+
+def test_ref_division():
+    ref = Ref(prefix="data", field="age")
+    math_op = ref / 21
+    assert math_op.to_str() == "(${data.age} / 21)"
+
+
+def test_ref_right_division():
+    ref = Ref(prefix="data", field="age")
+    math_op = 21 / ref
+    assert math_op.to_str() == "(21 / ${data.age})"
+
+
+def test_ref_modulus():
+    ref = Ref(prefix="data", field="age")
+    math_op = ref % 21
+    assert math_op.to_str() == "(${data.age} % 21)"
+
+
+def test_ref_right_modulus():
+    ref = Ref(prefix="data", field="age")
+    math_op = 21 % ref
+    assert math_op.to_str() == "(21 % ${data.age})"
+
+
+def test_math_ref_with_ref():
+    ref_one = Ref(prefix="data", field="age")
+    ref_two = Ref(prefix="data", field="height")
+    math_op = ref_one + ref_two
+    assert math_op.to_str() == "(${data.age} + ${data.height})"
+
+
+def test_ref_with_math_op():
+    ref = Ref(prefix="data", field="age")
+    math_op = 21 + ref
+    ref_with_math_op = ref + math_op
+    assert ref_with_math_op.to_str() == "(${data.age} + (21 + ${data.age}))"
+
+
+def test_math_op_with_ref():
+    ref = Ref(prefix="data", field="age")
+    math_op = 21 + ref
+    math_op_with_ref = math_op + ref
+    assert math_op_with_ref.to_str() == "((21 + ${data.age}) + ${data.age})"
+
+
+def test_ref_with_math_op_and_ref():
+    ref = Ref(prefix="data", field="age")
+    math_op = 21 + ref
+    math_op_with_ref = ref + math_op + ref
+    assert (
+        math_op_with_ref.to_str()
+        == "((${data.age} + (21 + ${data.age})) + ${data.age})"
+    )
+
+
 def test_logical_and_with_ref():
     ref1 = Ref(prefix="data", field="age")
     ref2 = Ref(prefix="form", field="is_verified")