Repository: PrefectHQ/prefect

Issue:
Make new conditional API compatible with the imperative API
## Current behavior
*Please describe how the feature works today*
In #2443, @jcrist added a new conditional API to Prefect, which is 💯. However, at this moment, it can only be used with the functional API (as it [requires](https://github.com/PrefectHQ/prefect/blob/12662a3a821b52dca0c0d39ced1bbbbba0674a85/src/prefect/tasks/control_flow/case.py#L87) a `flow` in the Prefect context). 


## Proposed behavior
*Please describe your proposed change to the current behavior*

The check for `case` in context currently [takes place](https://github.com/PrefectHQ/prefect/blob/12662a3a821b52dca0c0d39ced1bbbbba0674a85/src/prefect/core/task.py#L463) in `Task.bind()`, which is called by the functional API but not common imperative patterns like `Task.set_upstream()`, `Task.set_downstream()`, `Task.set_dependencies()`, `Flow.add_task()`, or `Flow.set_dependencies()`. 

I believe `Flow.add_task()` is the lowest-level API (called by all the others mentioned above) and consequently might be an opportune place to move the check for `case` in context. This way, adding a task to a flow while in a conditional block would automatically add it to the `case` class as well. Since the flow is (obviously) available when `Flow.add_task()` is called, it can be used inside the `case` context manager when calling `bind()` and `set_upstream()`. 

I think this would enable the following imperative syntax, while still using the existing case context managers:

```
Switch = Task()
A = Task()
B = Task()
C = Task()
f = Flow('imperative flow')

switch = Switch()
f.add_task(switch)

with case(switch, 'foo'):
    f.add_task(A())

with case(switch, 'bar'):
    b = B()
    c = C()
    c.set_upstream(b, flow=f)
```


## Example
*Please give an example of how the enhancement would be useful*



Discussion:
`case` needs a reference to the current flow to work properly, so I think this would have to be more like:

```python
with case(switch, 'bar', flow=f):
    ...
```
I was playing with that, but since `Case.add_task()` is always called from a parent method that has a `flow` object available, I think we can actually hide that from the user. For example, in the current implementation, it's called from `Task.bind()` which has a [flow object](https://github.com/PrefectHQ/prefect/blob/master/src/prefect/core/task.py#L457) in scope, so the signature could be `Case.add_task(task: Task, flow: Flow=None)`. We might want to validate that all passed flows are the same for a given case context manager, but by passing it to `add_task` we can probably hide the need for users to pass it explicitly

Resolving patch:
diff --git a/src/prefect/core/flow.py b/src/prefect/core/flow.py
--- a/src/prefect/core/flow.py
+++ b/src/prefect/core/flow.py
@@ -437,10 +437,13 @@ def add_task(self, task: Task) -> Task:
                     "flow.".format(task.slug)
                 )
 
-        if task not in self.tasks:
             self.tasks.add(task)
             self._cache.clear()
 
+            case = prefect.context.get("case", None)
+            if case is not None:
+                case.add_task(task, self)
+
         return task
 
     def add_edge(
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
@@ -461,10 +461,6 @@ 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/case.py b/src/prefect/tasks/control_flow/case.py
--- a/src/prefect/tasks/control_flow/case.py
+++ b/src/prefect/tasks/control_flow/case.py
@@ -1,7 +1,7 @@
 from typing import Any
 
 import prefect
-from prefect import Task
+from prefect import Task, Flow
 
 from .conditional import CompareValue
 
@@ -61,13 +61,21 @@ def __init__(self, task: Task, value: Any):
         self.task = task
         self.value = value
         self._tasks = set()
+        self._flow = None
 
-    def add_task(self, task: Task) -> None:
+    def add_task(self, task: Task, flow: Flow) -> None:
         """Add a new task under the case statement.
 
         Args:
             - task (Task): the task to add
+            - flow (Flow): the flow to use
         """
+        if self._flow is None:
+            self._flow = flow
+        elif self._flow is not flow:
+            raise ValueError(
+                "Multiple flows cannot be used with the same case statement"
+            )
         self._tasks.add(task)
 
     def __enter__(self):
@@ -81,17 +89,12 @@ def __exit__(self, *args):
             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
+                value=self.task, flow=self._flow
             )
 
             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)
+                if not self._tasks.intersection(self._flow.upstream_tasks(child)):
+                    child.set_upstream(cond, flow=self._flow)


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
@@ -464,3 +464,53 @@ def test_nested_case_execution(self, branch1, branch2):
                 assert state.result[t].result == sol[t]
             else:
                 assert state.result[t].is_skipped()
+
+    def test_case_imperative_errors(self):
+        flow = Flow("test")
+        flow2 = Flow("test2")
+        cond = identity.copy()
+        a = identity.copy()
+        with pytest.raises(ValueError, match="Multiple flows"):
+            with case(cond, True):
+                flow.add_task(a)
+                flow2.add_task(a)
+
+    @pytest.mark.parametrize("branch", ["a", "b", "c"])
+    def test_case_imperative_api(self, branch):
+        flow = Flow("test")
+
+        cond = identity.copy()
+        a = identity.copy()
+        b = inc.copy()
+        c = identity.copy()
+        d = inc.copy()
+
+        cond.bind(branch, flow=flow)
+        flow.add_task(cond)
+        with case(cond, "a"):
+            a.bind(1, flow=flow)
+            b.bind(a, flow=flow)
+            flow.add_task(a)
+            flow.add_task(b)
+
+        with case(cond, "b"):
+            c.bind(3, flow=flow)
+            d.bind(c, flow=flow)
+            flow.add_task(c)
+            flow.add_task(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()
+        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
+        elif branch == "c":
+            for t in [a, b, c, d]:
+                assert state.result[t].is_skipped()