Repository: PrefectHQ/prefect

Issue:
New conditional api
## Background

The current conditional api (`switch`, `ifelse`, and `merge`) work, but have a few problems. Using `ifelse` as an example:

```python
ifelse(cond : Task, true_task : Task, false_task : Task) -> None
```

- If `cond` evaluates to `True`, `true_task` is run and `false_task` is skipped. However, tasks upstream to `false_task` aren't skipped. For example:

  ```python
  ifelse(
      some_condition(),
      load(transform(extract())),
      some_other_task(),
  )
  ```

  if `some_condition` evaluates to `False`, only the `load` task is skipped, `transform` and `extract` still execute.

- If the result of a task is conditional based on upstream tasks, you'll need to call `merge` explicitly. We tried making `ifelse`/`switch` return the merge implicitly, but had to revert due to issues with implicitly creating terminal tasks in the graph (#2310, #2379).

  ```python
  on_true = do_something()
  on_false = do_something_else()
  cond = condition()
  ifelse(cond, on_true, on_false)
  result = merge(on_true, on_false)
  ```

  This approach works, but can be tricky to read.

## Proposal

I propose we add a contextmanager api for expressing conditional flows. The existing `switch`/`ifelse`/`merge` api will still exist (and may be used internally by the contextmanager api).

The api consists of two new functions:

- `case`: used to express branches in logic
- `var`: used to express merging outputs (this could equally be a class `Variable`)

A few examples to illustrate use:

### Example 1

Here we use `case` to express an `ifelse` branch:

```python
with Flow("example-1") as flow:
    cond = some_condition()
    
    with case(cond, True):
        a = do_a()
        b = do_b()
        do_c(a, b)

    with case(cond, False):
        do_c(1, 2)
```

If `some_condition` is `True`, `do_a`, `do_b`, and `do_c` are run. If it's false, `do_c` is run with some constant inputs. For each `case` block, a `Condition` task is created. If a task is created inside the `case` block *and* has no upstream dependencies that were also created inside that `case` block, then the condition is set as an upstream dependency. In this case, `a` and `b` would have a `Condition` for `True` set as an upstream dependency, but `do_c` would only depend on `a` and `b`. We could equally have all tasks created inside the block have the condition as an immediate upstream task, but that makes the graph visually complicated.

### Example 2

Here we illustrate merging tasks with `var`. Lets say we wanted to capture the output of `do_c` and use it later on. In normal Python code we might write this as:

```python
def example_2():
    if some_condition():
        a = do_a()
        b = do_b()
        c = do_c(a, b)
    else:
        c = do_c(1, 2)

    do_d(c)
```

Using this proposal, an equivalent flow would be:

```python
with Flow("example-2") as flow:
    cond = some_condition()
    
    c = var("c")

    with case(cond, True):
        a = do_a()
        b = do_b()
        c.set(do_c(a, b))

    with case(cond, False):
        c.set(do_c(1, 2))

    do_d(c)
```

Each call to `c.set` binds the value of `c` inside its corresponding `case` block. Behind the scenes this builds up a merge task branch-by-branch (note: we may not use merge directly, but the general structure will be the same). Since `c` is a `Task` itself, it can then be fed to downstream tasks.

### Example 3

These context managers should work fine even with nested conditional logic. For example, the following function:

```python
def example_3():
    if some_condition():
        if other_condition():
            a = do_a()
        else:
            a = 1
        b = do_b()
        do_c(a, b)
    else:
        d = do_d()
        do_e(d)
```

corresponds to an equivalent flow of:

```python
with Flow() as flow:
    cond = some_condition()

    with case(cond, True):

        cond2 = other_condition()

        a = var("a")
        with case(cond2, True):
            a.set(do_a())
        with case(cond2, False):
            a.set(1)

        b = do_b()
        do_c(a, b)

    with case(cond, False):
        d = do_d()
        do_e(d)
```

Note that in this case, the `Condition` task created by calling `case` on `cond2` will have the `Condition` created for `cond1` as a direct upstream task (following the logic described above, since it has no upstream tasks created inside this `case` context, the corresponding condition is set as a direct upstream task).

---

I think implementing the api as described above is fairly doable. It *might* require #2298 though to properly handle complex branching/merging logic (not sure yet).


Discussion:
Should this be more formalized as a PIN? Doesn't have to be very long and this seems like really cool functionality 😄 
Maybe? I'd be happy to write one up if that's the route we should go.
@jcrist just for posterity from talking in core release planning, the temperature of the group was to carry on with implementation without a dedicated PIN given that the new API is additive and does not change the existing API. If implementation does tread into breaking change territory for `iflese`/`switch`/`merge`, for example if there is any logic shared between the two APIs in the pipeline for branching that must be changed to add the new API, it sounded like we were of the opinion we should pursue a PIN in that case and futhermore determine a release schedule in core planning

Resolving patch:
diff --git a/src/prefect/core/task.py b/src/prefect/core/task.py
--- a/src/prefect/core/task.py
+++ b/src/prefect/core/task.py
@@ -436,6 +436,10 @@ def bind(
         if not flow:
             raise ValueError("Could not infer an active Flow context.")
 
+        case = prefect.context.get("case", None)
+        if case is not None:
+            case.add_task(self)
+
         self.set_dependencies(
             flow=flow,
             upstream_tasks=upstream_tasks,
diff --git a/src/prefect/tasks/control_flow/__init__.py b/src/prefect/tasks/control_flow/__init__.py
--- a/src/prefect/tasks/control_flow/__init__.py
+++ b/src/prefect/tasks/control_flow/__init__.py
@@ -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
diff --git a/src/prefect/tasks/control_flow/case.py b/src/prefect/tasks/control_flow/case.py
new file mode 100644
--- /dev/null
+++ b/src/prefect/tasks/control_flow/case.py
@@ -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)


Test patch:
diff --git a/tests/tasks/test_control_flow.py b/tests/tasks/test_control_flow.py
--- a/tests/tasks/test_control_flow.py
+++ b/tests/tasks/test_control_flow.py
@@ -4,7 +4,8 @@
 from prefect import Flow, Task, task
 from prefect.engine.result import NoResult
 from prefect.engine.state import Skipped, Success
-from prefect.tasks.control_flow import FilterTask, ifelse, merge, switch
+from prefect.tasks.control_flow import FilterTask, ifelse, merge, switch, case
+from prefect.tasks.control_flow.conditional import CompareValue
 from prefect.tasks.core.constants import Constant
 
 
@@ -324,3 +325,142 @@ def test_filter_func_can_be_changed(self):
         res = task.run([NoResult, NoResult, 0, 1, 5, "", exc])
         assert len(res) == 6
         assert res == [NoResult, NoResult, 0, 1, "", exc]
+
+
+@task
+def identity(x):
+    return x
+
+
+@task
+def inc(x):
+    return x + 1
+
+
+class TestCase:
+    def test_case_errors(self):
+        with Flow("test"):
+            with pytest.raises(TypeError, match="`value` cannot be a task"):
+                with case(identity(True), identity(False)):
+                    pass
+
+    def test_empty_case_block_no_tasks(self):
+        with Flow("test") as flow:
+            cond = identity(True)
+            with case(cond, True):
+                pass
+
+        # No tasks added if case block is empty
+        assert flow.tasks == {cond}
+
+    def test_case_sets_and_clears_context(self):
+        with Flow("test"):
+            c1 = case(identity(True), True)
+            c2 = case(identity(False), True)
+            assert "case" not in prefect.context
+            with c1:
+                assert prefect.context["case"] is c1
+                with c2:
+                    assert prefect.context["case"] is c2
+                assert prefect.context["case"] is c1
+            assert "case" not in prefect.context
+
+    def test_case_upstream_tasks(self):
+        with Flow("test") as flow:
+            a = identity(0)
+            with case(identity(True), True):
+                b = inc(a)
+                c = inc(b)
+                d = identity(1)
+                e = inc(d)
+                f = inc(e)
+
+        compare_values = [t for t in flow.tasks if isinstance(t, CompareValue)]
+        assert len(compare_values) == 1
+        comp = compare_values[0]
+
+        assert flow.upstream_tasks(a) == set()
+        assert flow.upstream_tasks(b) == {comp, a}
+        assert flow.upstream_tasks(c) == {b}
+        assert flow.upstream_tasks(d) == {comp}
+        assert flow.upstream_tasks(e) == {d}
+        assert flow.upstream_tasks(f) == {e}
+
+    @pytest.mark.parametrize("branch", ["a", "b", "c"])
+    def test_case_execution(self, branch):
+        with Flow("test") as flow:
+            cond = identity(branch)
+            with case(cond, "a"):
+                a = identity(1)
+                b = inc(a)
+
+            with case(cond, "b"):
+                c = identity(3)
+                d = inc(c)
+
+            e = merge(b, d)
+
+        state = flow.run()
+
+        if branch == "a":
+            assert state.result[a].result == 1
+            assert state.result[b].result == 2
+            assert state.result[c].is_skipped()
+            assert state.result[d].is_skipped()
+            assert state.result[e].result == 2
+        elif branch == "b":
+            assert state.result[a].is_skipped()
+            assert state.result[b].is_skipped()
+            assert state.result[c].result == 3
+            assert state.result[d].result == 4
+            assert state.result[e].result == 4
+        elif branch == "c":
+            for t in [a, b, c, d, e]:
+                assert state.result[t].is_skipped()
+
+    @pytest.mark.parametrize("branch1", [True, False])
+    @pytest.mark.parametrize("branch2", [True, False])
+    def test_nested_case_execution(self, branch1, branch2):
+        with Flow("test") as flow:
+            cond1 = identity(branch1)
+
+            a = identity(0)
+            with case(cond1, True):
+                cond2 = identity(branch2)
+                b = identity(10)
+                with case(cond2, True):
+                    c = inc(a)
+                    d = inc(c)
+                with case(cond2, False):
+                    e = inc(b)
+                    f = inc(e)
+
+                g = merge(d, f)
+
+            with case(cond1, False):
+                h = identity(3)
+                i = inc(h)
+
+            j = merge(g, i)
+
+        state = flow.run()
+
+        sol = {a: 0, cond1: branch1}
+        if branch1:
+            sol[cond2] = branch2
+            sol[b] = 10
+            if branch2:
+                sol[c] = 1
+                sol[d] = sol[g] = sol[j] = 2
+            else:
+                sol[e] = 11
+                sol[f] = sol[g] = sol[j] = 12
+        else:
+            sol[h] = 3
+            sol[i] = sol[j] = 4
+
+        for t in [cond1, cond2, a, b, c, d, e, f, g, h, i, j]:
+            if t in sol:
+                assert state.result[t].result == sol[t]
+            else:
+                assert state.result[t].is_skipped()