-
Notifications
You must be signed in to change notification settings - Fork 1.5k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #2443 from jcrist/new-conditional-api
New conditional api
- Loading branch information
Showing
7 changed files
with
277 additions
and
10 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,35 +1,60 @@ | ||
# Control Flow | ||
|
||
Tasks and utilities for implementing control flow constructs like branching and rejoining flows. | ||
Tasks and utilities for implementing control flow constructs like branching and | ||
rejoining flows. | ||
|
||
## Case <Badge text="fn"/> | ||
|
||
A conditional block in a flow. | ||
|
||
Used as a context-manager, ``case(task, value)`` creates a block of tasks that | ||
are only run if the result of ``task`` is equal to ``value``. | ||
|
||
[API Reference](/api/latest/tasks/control_flow.html#prefect-tasks-control-flow-case-case) | ||
|
||
## If/Else <Badge text="fn"/> | ||
|
||
Builds a conditional branch into a workflow. | ||
|
||
If the condition evaluates True(ish), the true_task will run. If it evaluates False(ish), the false_task will run. The task doesn't run is Skipped, as are all downstream tasks that don't set `skip_on_upstream_skip=False`. | ||
If the condition evaluates True(ish), the ``true_task`` will run. If it | ||
evaluates False(ish), the ``false_task`` will run. The task that doesn't run is | ||
Skipped, as are all downstream tasks that don't set | ||
`skip_on_upstream_skip=False`. | ||
|
||
[API Reference](/api/latest/tasks/control_flow.html#prefect-tasks-control-flow-conditional-ifelse) | ||
|
||
|
||
## Switch <Badge text="task"/> | ||
## Switch <Badge text="fn"/> | ||
|
||
Adds a SWITCH to a workflow. | ||
|
||
The condition task is evaluated and the result is compared to the keys of the cases dictionary. The task corresponding to the matching key is run; all other tasks are skipped. Any tasks downstream of the skipped tasks are also skipped unless they set `skip_on_upstream_skip=False`. | ||
The condition task is evaluated and the result is compared to the keys of the | ||
cases dictionary. The task corresponding to the matching key is run; all other | ||
tasks are skipped. Any tasks downstream of the skipped tasks are also skipped | ||
unless they set `skip_on_upstream_skip=False`. | ||
|
||
[API Reference](/api/latest/tasks/control_flow.html#prefect-tasks-control-flow-conditional-switch) | ||
|
||
|
||
## Merge <Badge text="task"/> | ||
|
||
Merges conditional branches back together. | ||
|
||
A conditional branch in a flow results in one or more tasks proceeding and one or more tasks skipping. It is often convenient to merge those branches back into a single result. This function is a simple way to achieve that goal. | ||
A conditional branch in a flow results in one or more tasks proceeding and one | ||
or more tasks skipping. It is often convenient to merge those branches back | ||
into a single result. This function is a simple way to achieve that goal. | ||
|
||
The merge will return the first real result it encounters, or `None`. If multiple tasks might return a result, group them with a list. | ||
The merge will return the first real result it encounters, or `None`. If | ||
multiple tasks might return a result, group them with a list. | ||
|
||
[API Reference](/api/latest/tasks/control_flow.html#prefect-tasks-control-flow-conditional-merge) | ||
|
||
## FilterTask <Badge text="task"/> | ||
|
||
Task for filtering lists of results. | ||
|
||
The default filter is to filter out `NoResults` and `Exceptions` for filtering out mapped results. Note that this task has a default trigger of `all_finished` and `skip_on_upstream_skip=False`. | ||
The default filter is to filter out `NoResults` and `Exceptions` for filtering | ||
out mapped results. Note that this task has a default trigger of `all_finished` | ||
and `skip_on_upstream_skip=False`. | ||
|
||
[API Reference](/api/latest/tasks/control_flow.html#prefect-tasks-control-flow-filter-filtertask) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,2 +1,3 @@ | ||
from prefect.tasks.control_flow.conditional import ifelse, switch, merge | ||
from prefect.tasks.control_flow.filter import FilterTask | ||
from prefect.tasks.control_flow.case import case |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,97 @@ | ||
from typing import Any | ||
|
||
import prefect | ||
from prefect import Task | ||
|
||
from .conditional import CompareValue | ||
|
||
|
||
__all__ = ("case",) | ||
|
||
|
||
class case(object): | ||
"""A conditional block in a flow definition. | ||
Used as a context-manager, ``case`` creates a block of tasks that are only | ||
run if the result of ``task`` is equal to ``value``. | ||
Args: | ||
- task (Task): The task to use in the comparison | ||
- value (Any): A constant the result of ``task`` will be compared with | ||
Example: | ||
A ``case`` block is similar to Python's if-blocks. It delimits a block | ||
of tasks that will only be run if the result of ``task`` is equal to | ||
``value``: | ||
```python | ||
# Standard python code | ||
if task == value: | ||
res = run_if_task_equals_value() | ||
other_task(res) | ||
# Equivalent prefect code | ||
with case(task, value): | ||
# Tasks created in this block are only run if the | ||
# result of ``task`` is equal to ``value`` | ||
res = run_if_task_equals_value() | ||
other_task(run) | ||
``` | ||
The ``value`` argument can be any non-task object. Here we branch on a | ||
string result: | ||
```python | ||
with Flow("example") as flow: | ||
cond = condition() | ||
with case(cond, "a"): | ||
run_if_cond_is_a() | ||
with case(cond, "b"): | ||
run_if_cond_is_b() | ||
``` | ||
""" | ||
|
||
def __init__(self, task: Task, value: Any): | ||
if isinstance(value, Task): | ||
raise TypeError("`value` cannot be a task") | ||
|
||
self.task = task | ||
self.value = value | ||
self._tasks = set() | ||
|
||
def add_task(self, task: Task) -> None: | ||
"""Add a new task under the case statement. | ||
Args: | ||
- task (Task): the task to add | ||
""" | ||
self._tasks.add(task) | ||
|
||
def __enter__(self): | ||
self.__prev_case = prefect.context.get("case") | ||
prefect.context.update(case=self) | ||
|
||
def __exit__(self, *args): | ||
if self.__prev_case is None: | ||
prefect.context.pop("case", None) | ||
else: | ||
prefect.context.update(case=self.__prev_case) | ||
|
||
if self._tasks: | ||
|
||
flow = prefect.context.get("flow", None) | ||
if not flow: | ||
raise ValueError("Could not infer an active Flow context.") | ||
|
||
cond = CompareValue(self.value, name=f"case({self.value})",).bind( | ||
value=self.task | ||
) | ||
|
||
for child in self._tasks: | ||
# If a task has no upstream tasks created in this case block, | ||
# the case conditional should be set as an upstream task. | ||
if not self._tasks.intersection(flow.upstream_tasks(child)): | ||
child.set_upstream(cond) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters