Repository: PrefectHQ/prefect

Issue:
Chain `set_upstream` with Task Results
## Current behavior
<!-- Please describe how the feature works today -->
The `set_upstream` method is available on the functional API for tasks to explicitly set upstream dependencies like so:
```
from prefect import task, Flow


@task
def foo():
    return "foo"


@task
def bar():
    return "bar"


@task
def oops(x):
    return "oops"


with Flow("test") as flow:
    bar_result = bar()
    foo_result = foo()
    foo_result.set_upstream(bar_result)
    oops_result = oops(foo_result)
```



## Proposed behavior
<!-- Please describe your proposed change to the current behavior -->
It would be intuitive to use the API with chaining, like so:

```
from prefect import task, Flow


@task
def foo():
    return "foo"


@task
def bar():
    return "bar"


@task
def oops(x):
    return "oops"


with Flow("test") as flow:
    bar_result = bar()
    foo_result = foo().set_upstream(bar_result)
    oops_results = oops(foo_result)
```



Discussion:

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
@@ -635,7 +635,7 @@ def set_dependencies(
         keyword_tasks: Mapping[str, object] = None,
         mapped: bool = False,
         validate: bool = None,
-    ) -> None:
+    ) -> "Task":
         """
         Set dependencies for a flow either specified or in the current context using this task
 
@@ -653,7 +653,7 @@ def set_dependencies(
                 configuration file.
 
         Returns:
-            - None
+            - self
 
         Raises:
             - ValueError: if no flow is specified and no flow can be found in the current context
@@ -673,9 +673,11 @@ def set_dependencies(
             mapped=mapped,
         )
 
+        return self
+
     def set_upstream(
         self, task: object, flow: "Flow" = None, key: str = None, mapped: bool = False
-    ) -> None:
+    ) -> "Task":
         """
         Sets the provided task as an upstream dependency of this task.
 
@@ -690,7 +692,7 @@ def set_upstream(
             - mapped (bool, optional): Whether this dependency is mapped; defaults to `False`
 
         Returns:
-            - None
+            - self
 
         Raises:
             - ValueError: if no flow is specified and no flow can be found in the current context
@@ -701,9 +703,11 @@ def set_upstream(
         else:
             self.set_dependencies(flow=flow, upstream_tasks=[task], mapped=mapped)
 
+        return self
+
     def set_downstream(
         self, task: "Task", flow: "Flow" = None, key: str = None, mapped: bool = False
-    ) -> None:
+    ) -> "Task":
         """
         Sets the provided task as a downstream dependency of this task.
 
@@ -715,6 +719,9 @@ def set_downstream(
                 will be passed to the downstream task's `run()` method under this keyword argument.
             - mapped (bool, optional): Whether this dependency is mapped; defaults to `False`
 
+        Returns:
+            - self
+
         Raises:
             - ValueError: if no flow is specified and no flow can be found in the current context
         """
@@ -726,6 +733,8 @@ def set_downstream(
         else:
             task.set_dependencies(flow=flow, upstream_tasks=[self], mapped=mapped)
 
+        return self
+
     def inputs(self) -> Dict[str, Dict]:
         """
         Describe the inputs for this task. The result is a dictionary that maps each input to


Test patch:
diff --git a/tests/core/test_task.py b/tests/core/test_task.py
--- a/tests/core/test_task.py
+++ b/tests/core/test_task.py
@@ -529,6 +529,21 @@ def test_set_upstream_with_properties(self, props):
             t2.set_upstream(t1, **props)
             assert Edge(t1, t2, **props) in f.edges
 
+    def test_set_dependencies_stream_allows_chaining(self):
+        t1 = Task()
+        t2 = Task()
+        t3 = Task()
+        with Flow(name="test") as f:
+            t1_result = t1()
+            t2_result = t2()
+            t3_result = t3()
+
+            assert t1_result.set_downstream(t2_result) is t1_result
+            assert t3_result.set_upstream(t2_result) is t3_result
+            assert (
+                t3_result.set_dependencies(f, upstream_tasks=[t1_result]) is t3_result
+            )
+
 
 class TestSerialization:
     def test_serialization(self):