Skip to content

Commit

Permalink
Adds higher-level conditions functions
Browse files Browse the repository at this point in the history
cond_1 and cond_2
cond_1 or cond_2
  • Loading branch information
elijahbenizzy committed May 22, 2024
1 parent 30aac6c commit 7735ada
Show file tree
Hide file tree
Showing 3 changed files with 78 additions and 1 deletion.
54 changes: 53 additions & 1 deletion burr/core/action.py
Original file line number Diff line number Diff line change
Expand Up @@ -182,7 +182,11 @@ class Condition(Function):
def __init__(self, keys: List[str], resolver: Callable[[State], bool], name: str = None):
"""Base condition class. Chooses keys to read from the state and a resolver function.
Note that the ``~`` operator allows you to automatically invert the condition.
Note that you can use a few fundamental operators to build more complex conditions:
- ``~`` operator allows you to automatically invert the condition.
- ``|`` operator allows you to OR two conditions together.
- ``&`` operator allows you to AND two conditions together.
:param keys: Keys to read from the state
:param resolver: Function to resolve the condition to True or False
Expand Down Expand Up @@ -281,6 +285,54 @@ def default(self) -> "Condition":
def name(self) -> str:
return self._name

def __or__(self, other: "Condition") -> "Condition":
"""Combines two conditions with an OR operator. This will return a new condition
that is the OR of the two conditions.
To check if either foo is bar or baz is qux:
.. code-block:: python
condition = Condition.when(foo="bar") | Condition.when(baz="qux")
:param other: Other condition to OR with
:return: A new condition that is the OR of the two conditions
"""
if not isinstance(other, Condition):
raise ValueError(f"Cannot OR a Condition with {other}")
return Condition(
self._keys + other._keys,
lambda state: self._resolver(state) or other.resolver(state),
name=f"{self._name} | {other._name}",
)

def __and__(self, other: "Condition") -> "Condition":
"""Combines two conditions with an AND operator. This will return a new condition
that is the AND of the two conditions.
To check if both foo is bar and baz is qux:
.. code-block:: python
condition = Condition.when(foo="bar") & Condition.when(baz="qux")
# equivalent to
condition = Condition.when(foo="bar", baz="qux")
:param other: Other condition to AND with
:return: A new condition that is the AND of the two conditions
"""
if not isinstance(other, Condition):
raise ValueError(f"Cannot AND a Condition with {other}")
return Condition(
self._keys + other._keys,
lambda state: self._resolver(state) and other.resolver(state),
name=f"{self._name} & {other._name}",
)

@property
def resolver(self) -> Callable[[State], bool]:
return self._resolver

def __invert__(self):
return Condition(self._keys, lambda state: not self._resolver(state), name=f"~{self._name}")

Expand Down
1 change: 1 addition & 0 deletions docs/reference/conditions.rst
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@ Conditions represent choices to move between actions -- these are read by the ap
Note that these will always be specified in order -- the first condition that evaluates to ``True`` will be the selected action.

.. autoclass:: burr.core.action.Condition
:special-members: __and__, __or__, __invert__
:members:

.. automethod:: __init__
24 changes: 24 additions & 0 deletions tests/core/test_action.py
Original file line number Diff line number Diff line change
Expand Up @@ -153,6 +153,30 @@ def test_condition_invert():
assert cond_inverted.run(State({"foo": "baz"})) == {Condition.KEY: True}


def test_condition_and():
cond1 = Condition.when(foo="bar")
cond2 = Condition.when(baz="qux")
cond_and = cond1 & cond2
assert cond_and.name == "foo=bar & baz=qux"
assert sorted(cond_and.reads) == ["baz", "foo"]
assert cond_and.run(State({"foo": "bar", "baz": "qux"})) == {Condition.KEY: True}
assert cond_and.run(State({"foo": "baz", "baz": "qux"})) == {Condition.KEY: False}
assert cond_and.run(State({"foo": "bar", "baz": "corge"})) == {Condition.KEY: False}
assert cond_and.run(State({"foo": "baz", "baz": "corge"})) == {Condition.KEY: False}


def test_condition_or():
cond1 = Condition.when(foo="bar")
cond2 = Condition.when(baz="qux")
cond_or = cond1 | cond2
assert cond_or.name == "foo=bar | baz=qux"
assert sorted(cond_or.reads) == ["baz", "foo"]
assert cond_or.run(State({"foo": "bar", "baz": "qux"})) == {Condition.KEY: True}
assert cond_or.run(State({"foo": "baz", "baz": "qux"})) == {Condition.KEY: True}
assert cond_or.run(State({"foo": "bar", "baz": "corge"})) == {Condition.KEY: True}
assert cond_or.run(State({"foo": "baz", "baz": "corge"})) == {Condition.KEY: False}


def test_result():
result = Result("foo", "bar")
assert result.run(State({"foo": "baz", "bar": "qux", "baz": "quux"})) == {
Expand Down

0 comments on commit 7735ada

Please sign in to comment.