Skip to content

Commit

Permalink
[flows] using math operators between math objs
Browse files Browse the repository at this point in the history
  • Loading branch information
@david-lev
david-lev committed Nov 28, 2024
1 parent db21a76 commit 2e8033b
Show file tree
Hide file tree
Showing 2 changed files with 125 additions and 82 deletions.
191 changes: 109 additions & 82 deletions pywa/types/flows.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1338,7 +1338,7 @@ def __post_init__(self):

def __truediv__(self, ref: _RefT) -> _RefT:
"""A shortcut to reference screen data / form components in this screen."""
return ref.__class__(ref.field, screen=self.id)
return ref.__class__(ref._field, screen=self.id)


class LayoutType(utils.StrEnum):
Expand Down Expand Up @@ -1414,99 +1414,141 @@ class ComponentType(utils.StrEnum):
SWITCH = "Switch"


class Ref:
"""Base class for all references"""
class _Expr:
"""Base for refs, conditions, and expressions"""

def __init__(self, prefix: str, field: str, screen: Screen | str | None = None):
self.prefix = prefix
self.field = field
self.screen = (
f"screen.{screen.id if isinstance(screen, Screen) else screen}."
if screen
else ""
)

def to_str(self) -> str:
return "${%s%s.%s}" % (
self.screen,
self.prefix,
self.field,
)
def to_str(self) -> str: ...

def __str__(self) -> str:
"""Allowing to use in string concatenation. Added in v6.0."""
return self.to_str()

@staticmethod
def _format_value(val: Ref | bool | int | float | str) -> str:
if isinstance(val, (Ref, MathExpression)):
def _format_value(val: _Expr | bool | int | float | str) -> str:
if isinstance(val, _Expr):
return val.to_str()
elif isinstance(val, str):
return f"'{val}'"
elif isinstance(val, bool):
return str(val).lower()
return str(val)

def __eq__(self, other: Ref | bool | int | float | str) -> Condition:
return Condition(f"({self.to_str()} == {self._format_value(other)})")

def __ne__(self, other: Ref | bool | int | float | str) -> Condition:
return Condition(f"({self.to_str()} != {self._format_value(other)})")

def __gt__(self, other: Ref | int | float) -> Condition:
return Condition(f"({self.to_str()} > {self._format_value(other)})")

def __ge__(self, other: Ref | int | float) -> Condition:
return Condition(f"({self.to_str()} >= {self._format_value(other)})")

def __lt__(self, other: Ref | int | float) -> Condition:
return Condition(f"({self.to_str()} < {self._format_value(other)})")
class _Math(_Expr):
"""Base for math expressions"""

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:
def __add__(
self: Ref | MathExpression, other: Ref | MathExpression | int | float
) -> MathExpression:
return MathExpression(f"({self.to_str()} + {self._format_value(other)})")

def __radd__(self, other: Ref | int | float) -> MathExpression:
def __radd__(
self: Ref | MathExpression, other: Ref | MathExpression | int | float
) -> MathExpression:
return MathExpression(f"({self._format_value(other)} + {self.to_str()})")

def __sub__(self, other: Ref | int | float) -> MathExpression:
def __sub__(
self: Ref | MathExpression, other: Ref | MathExpression | int | float
) -> MathExpression:
return MathExpression(f"({self.to_str()} - {self._format_value(other)})")

def __rsub__(self, other: Ref | int | float) -> MathExpression:
def __rsub__(
self: Ref | MathExpression, other: Ref | MathExpression | int | float
) -> MathExpression:
return MathExpression(f"({self._format_value(other)} - {self.to_str()})")

def __mul__(self, other: Ref | int | float) -> MathExpression:
def __mul__(
self: Ref | MathExpression, other: Ref | MathExpression | int | float
) -> MathExpression:
return MathExpression(f"({self.to_str()} * {self._format_value(other)})")

def __rmul__(self, other: Ref | int | float) -> MathExpression:
def __rmul__(
self: Ref | MathExpression, other: Ref | MathExpression | int | float
) -> MathExpression:
return MathExpression(f"({self._format_value(other)} * {self.to_str()})")

def __truediv__(self, other: Ref | int | float) -> MathExpression:
def __truediv__(
self: Ref | MathExpression, other: Ref | MathExpression | int | float
) -> MathExpression:
return MathExpression(f"({self.to_str()} / {self._format_value(other)})")

def __rtruediv__(self, other: Ref | int | float) -> MathExpression:
def __rtruediv__(
self: Ref | MathExpression, other: Ref | MathExpression | int | float
) -> MathExpression:
return MathExpression(f"({self._format_value(other)} / {self.to_str()})")

def __mod__(self, other: Ref | int | float) -> MathExpression:
def __mod__(
self: Ref | MathExpression, other: Ref | MathExpression | int | float
) -> MathExpression:
return MathExpression(f"({self.to_str()} % {self._format_value(other)})")

def __rmod__(self, other: Ref | int | float) -> MathExpression:
def __rmod__(
self: Ref | MathExpression, other: Ref | MathExpression | 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()})")
return Condition(f"({self.to_str()} && {other._expression})")

def __or__(self, other: Ref | Condition) -> Condition:
if isinstance(other, Ref):
return Condition(f"({self.to_str()} || {other.to_str()})")
return Condition(f"({self.to_str()} || {other._expression})")
class _Combine(_Expr):
""" "Base for combining refs and conditions"""

def _get_left_right(
self: Ref | Condition, right: Ref | Condition
) -> tuple[str, str]:
return self.to_str() if isinstance(
self, Ref
) else self._expression, right.to_str() if isinstance(
right, Ref
) else right._expression

def __invert__(self) -> Condition:
return Condition(f"!{self.to_str()}")
def __and__(self: Ref | Condition, other: Ref | Condition) -> Condition:
left, right = self._get_left_right(other)
return Condition(f"({left} && {right})")

def __or__(self: Ref | Condition, other: Ref | Condition) -> Condition:
left, right = self._get_left_right(other)
return Condition(f"({left} || {right})")

def __invert__(self: Ref | Condition) -> Condition:
return Condition(
f"!{self.to_str() if isinstance(self, Ref) else self._expression}"
)


class Ref(_Math, _Combine):
"""Base class for all references"""

def __init__(self, prefix: str, field: str, screen: Screen | str | None = None):
self._prefix = prefix
self._field = field
self._screen = (
f"screen.{screen.id if isinstance(screen, Screen) else screen}."
if screen
else ""
)

def to_str(self) -> str:
return "${%s%s.%s}" % (
self._screen,
self._prefix,
self._field,
)

def __eq__(self, other: Ref | bool | int | float | str) -> Condition:
return Condition(f"({self.to_str()} == {self._format_value(other)})")

def __ne__(self, other: Ref | bool | int | float | str) -> Condition:
return Condition(f"({self.to_str()} != {self._format_value(other)})")

def __gt__(self, other: Ref | int | float) -> Condition:
return Condition(f"({self.to_str()} > {self._format_value(other)})")

def __ge__(self, other: Ref | int | float) -> Condition:
return Condition(f"({self.to_str()} >= {self._format_value(other)})")

def __lt__(self, other: Ref | int | float) -> Condition:
return Condition(f"({self.to_str()} < {self._format_value(other)})")

def __le__(self, other: Ref | int | float) -> Condition:
return Condition(f"({self.to_str()} <= {self._format_value(other)})")

def __repr__(self) -> str:
return f"{self.__class__.__name__}({self.to_str()})"
Expand All @@ -1515,7 +1557,7 @@ def __repr__(self) -> str:
_RefT = TypeVar("_RefT", bound=Ref)


class MathExpression:
class MathExpression(_Math):
"""
This class automatically created when using the arithmetic operators on :class:`Ref` objects.
Expand Down Expand Up @@ -1565,17 +1607,14 @@ 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:
class Condition(_Combine):
"""
This class automatically created when using the comparison operators on :class:`Ref` objects.
Expand Down Expand Up @@ -1642,22 +1681,6 @@ def __init__(self, expression: str):
self._expression = expression
self.wrap_with_backticks = False

def __and__(self, other: Condition | Ref) -> Condition:
if isinstance(other, Condition):
return Condition(f"({self._expression} && {other._expression})")
return Condition(f"({self._expression} && {other.to_str()})")

def __or__(self, other: Condition | Ref) -> Condition:
if isinstance(other, Condition):
return Condition(f"({self._expression} || {other._expression})")
return Condition(f"({self._expression} || {other.to_str()})")

def __invert__(self) -> Condition:
return Condition(f"!{self._expression}")

def __str__(self) -> str:
return self.to_str()

def __repr__(self) -> str:
return f"Condition({self._expression})"

Expand All @@ -1671,12 +1694,12 @@ def to_str(self) -> str:

class ScreenDataRef(Ref):
"""
Represents a ScreenData reference (converts to ``${data.<key>}`` | ``${screen.<screen>.data.<key>}``).
Represents a :class:`ScreenData` reference (converts to ``${data.<key>}`` | ``${screen.<screen>.data.<key>}``).
Example:
- Hint: use this class directly only if you don't have access to the :class:`ScreenData` object.
- Hint: use the ``.ref`` property of :class:`ScreenData` to get reference to a ScreenData.
- Hint: use the ``.ref_in(screen)`` method of :class:`ScreenData` to get the data ref from another screen.
- Hint: use the ``.ref_in(screen)`` method of :class:`ScreenData` to get the data ref from another screen (or ``screen/ref``).
>>> FlowJSON(
... screens=[
Expand All @@ -1688,7 +1711,9 @@ class ScreenDataRef(Ref):
... TextHeading(
... text=welcome.ref, # data in the same screen
... visible=is_visible.ref_in(other) # data from other screen
... )
... ),
... TextBody(
... text=ScreenDataRef(key='welcome', screen='START'), # using the class directly
... ])
... )
... ]
Expand All @@ -1709,6 +1734,7 @@ class ComponentRef(Ref):
Example:
- Hint: use this class directly only if you don't have access to the component object.
- Hint: use the ``.ref`` property of each component to get a reference to this component.
- Hint: use the ``.ref_in(screen)`` method of each component to get the component reference variable of that component with the given screen name.
Expand All @@ -1724,6 +1750,7 @@ class ComponentRef(Ref):
... phone := TextInput(name='phone', ...),
... TextBody(text=phone.ref, ...), # component reference from the same screen
... TextCaption(text=email.ref_in(other), ...) # component reference from another screen
... TextHeading(text=ComponentRef('phone', screen='START'), ...) # using the class directly
... ])
... )
... ]
Expand Down
16 changes: 16 additions & 0 deletions tests/test_flows.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -239,6 +239,22 @@ def test_math_op_with_ref():
assert math_op_with_ref.to_str() == "((21 + ${data.age}) + ${data.age})"


def test_math_op_with_math_op():
ref = Ref(prefix="data", field="age")
math_op_1 = 21 + ref
math_op_2 = 21 + ref
assert (
math_op_1 + math_op_2
).to_str() == "((21 + ${data.age}) + (21 + ${data.age}))"


def test_math_op_with_number():
ref = Ref(prefix="data", field="age")
math_op = 21 + ref
math_op_with_number = math_op + 21
assert math_op_with_number.to_str() == "((21 + ${data.age}) + 21)"


def test_ref_with_math_op_and_ref():
ref = Ref(prefix="data", field="age")
math_op = 21 + ref
Expand Down

0 comments on commit 2e8033b

Please sign in to comment.