Repository: PrefectHQ/prefect

Issue:
`ifelse` should behave like a ternary expression
Currently, our conditional `ifelse` operator just sets dependencies in a particular way, but always returns `None`.  It would feel much more natural if it returned a task object whose value at runtime behaved like a standard `x = condition ? if_true : if_false` expression


Discussion:

Resolving patch:
diff --git a/src/prefect/tasks/control_flow/conditional.py b/src/prefect/tasks/control_flow/conditional.py
--- a/src/prefect/tasks/control_flow/conditional.py
+++ b/src/prefect/tasks/control_flow/conditional.py
@@ -82,11 +82,16 @@ def b_branch():
             The value of the `condition` task will be compared to the keys of this dict, and
             the matching task will be executed.
 
+    Returns:
+        - Task: a task whose result is the output from the task executed by this switch
+
     Raises:
         - PrefectWarning: if any of the tasks in "cases" have upstream dependencies,
-            then this task will warn that those upstream tasks may run whether or not the switch condition matches their branch. The most common cause of this
-            is passing a list of tasks as one of the cases, which adds the `List` task
-            to the switch condition but leaves the tasks themselves upstream.
+            then this task will warn that those upstream tasks may run whether
+            or not the switch condition matches their branch. The most common
+            cause of this is passing a list of tasks as one of the cases, which
+            adds the `List` task to the switch condition but leaves the tasks
+            themselves upstream.
     """
 
     with prefect.tags("switch"):
@@ -94,6 +99,7 @@ def b_branch():
             task = prefect.utilities.tasks.as_task(task)
             match_condition = CompareValue(value=value).bind(value=condition)
             task.set_dependencies(upstream_tasks=[match_condition])
+        return merge(*cases.values())
 
 
 def ifelse(condition: Task, true_task: Task, false_task: Task) -> None:
@@ -101,20 +107,26 @@ def ifelse(condition: Task, true_task: Task, false_task: Task) -> None:
     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`.
+    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`.
 
     Args:
         - condition (Task): a task whose boolean result forms the condition for the ifelse
         - true_task (Task): a task that will be executed if the condition is True
         - false_task (Task): a task that will be executed if the condition is False
+
+    Returns:
+        - Task: a task whose result is the output from the task executed by this ifelse
     """
 
     @prefect.task
     def as_bool(x):
         return bool(x)
 
-    switch(condition=as_bool(condition), cases={True: true_task, False: false_task})
+    return switch(
+        condition=as_bool(condition), cases={True: true_task, False: false_task}
+    )
 
 
 def merge(*tasks: Task) -> Task:


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
@@ -5,7 +5,7 @@
 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.core.constants import Constant
+from prefect.utilities.tasks import as_task
 
 
 class Condition(Task):
@@ -18,25 +18,32 @@ def run(self):
         return 1
 
 
+@task
+def identity(x):
+    return x
+
+
 @pytest.mark.parametrize("condition_value", [True, False])
 def test_ifelse(condition_value):
     condition = Condition()
-    true_branch = SuccessTask(name="true branch")
-    false_branch = SuccessTask(name="false branch")
 
     with Flow(name="test") as flow:
-        cnd = ifelse(condition, true_branch, false_branch)
-        assert len(flow.tasks) == 6
+        true_branch = identity("true")
+        false_branch = identity("false")
+        res = ifelse(condition, true_branch, false_branch)
+        assert len(flow.tasks) == 7
 
     with prefect.context(CONDITION=condition_value):
         state = flow.run()
 
-    assert isinstance(
-        state.result[true_branch], Success if condition_value is True else Skipped
-    )
-    assert isinstance(
-        state.result[false_branch], Success if condition_value is False else Skipped
-    )
+    if condition_value:
+        assert state.result[true_branch].is_successful()
+        assert state.result[false_branch].is_skipped()
+        assert state.result[res].result == "true"
+    else:
+        assert state.result[true_branch].is_skipped()
+        assert state.result[false_branch].is_successful()
+        assert state.result[res].result == "false"
 
 
 @pytest.mark.parametrize("condition_value", [1, "a", "False", True, [1], {1: 2}])
@@ -46,7 +53,7 @@ def test_ifelse_with_truthy_conditions(condition_value):
     false_branch = SuccessTask(name="false branch")
 
     with Flow(name="test") as flow:
-        cnd = ifelse(condition, true_branch, false_branch)
+        ifelse(condition, true_branch, false_branch)
 
     with prefect.context(CONDITION=condition_value):
         state = flow.run()
@@ -62,7 +69,7 @@ def test_ifelse_with_falsey_conditions(condition_value):
     false_branch = SuccessTask(name="false branch")
 
     with Flow(name="test") as flow:
-        cnd = ifelse(condition, true_branch, false_branch)
+        ifelse(condition, true_branch, false_branch)
 
     with prefect.context(CONDITION=condition_value):
         state = flow.run()
@@ -74,14 +81,15 @@ def test_ifelse_with_falsey_conditions(condition_value):
 @pytest.mark.parametrize("condition_value", ["a", "b", "c", "d", "x"])
 def test_switch(condition_value):
     condition = Condition()
-    a_branch = SuccessTask(name="a")
-    b_branch = SuccessTask(name="b")
-    c_branch = SuccessTask(name="c")
-    d_branch = SuccessTask(name="d")
 
     with Flow(name="test") as flow:
-        switch(condition, dict(a=a_branch, b=b_branch, c=c_branch, d=d_branch))
-        assert len(flow.tasks) == 9
+        a_branch = identity("a_val")
+        b_branch = identity("b_val")
+        c_branch = identity("c_val")
+        d_branch = identity("d_val")
+
+        res = switch(condition, dict(a=a_branch, b=b_branch, c=c_branch, d=d_branch))
+        assert len(flow.tasks) == 10
 
     with prefect.context(CONDITION=condition_value):
         state = flow.run()
@@ -97,6 +105,10 @@ def test_switch(condition_value):
         assert isinstance(
             state.result[d_branch], Success if condition_value == "d" else Skipped
         )
+        if condition_value == "x":
+            assert isinstance(state.result[res], Skipped)
+        else:
+            assert state.result[res].result == f"{condition_value}_val"
 
 
 def test_merging_diamond_flow():
@@ -172,8 +184,9 @@ def test_list_of_tasks():
     with Flow(name="test") as flow:
         condition = Condition()
         true_branch = [SuccessTask(), SuccessTask()]
+        list_task = as_task(true_branch)
         false_branch = SuccessTask()
-        ifelse(condition, true_branch, false_branch)
+        ifelse(condition, list_task, false_branch)
 
     with prefect.context(CONDITION=True):
         state = flow.run()
@@ -188,9 +201,6 @@ def test_list_of_tasks():
         for t in true_branch:
             # the tasks in the list ran becuase they have no upstream dependencies.
             assert isinstance(state.result[t], Success)
-        list_task = next(
-            t for t in flow.tasks if isinstance(t, prefect.tasks.core.collections.List)
-        )
         # but the list itself skipped
         assert isinstance(state.result[list_task], Skipped)
         assert isinstance(state.result[false_branch], Success)