Repository: PrefectHQ/prefect

Issue:
task.set_*stream should have more keyword arguments
`Task.set_upstream` and `Task.set_downstream` were originally envisioned as methods for building up a Flow which does not pass data.  `Task.set_dependencies` offers a richer API for specifying data passage between tasks.

We should elevate all these methods to expose the full set of Prefect's features:
- allow for the creation of keyed edges
- allow for the creation of mapped edges

We might need to think about whether `set_upstream` and `set_downstream` should have the same exact options - calls where the upstream task specifies a key on the downstream task like `t1.set_downstream(t2, key="y")` feel unnatural.


Discussion:

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
@@ -472,9 +472,8 @@ def add_edge(
         Args:
             - upstream_task (Task): The task that the edge should start from
             - downstream_task (Task): The task that the edge should end with
-            - key (str, optional): The key to be set for the new edge; this is
-                the argument name the result of the upstream task will be bound to in the
-                `run()` method of the downstream task
+            - key (str, optional): The key to be set for the new edge; the result of the upstream task
+            will be passed to the downstream task's `run()` method under this keyword argument
             - mapped (bool, optional): Whether this edge represents a call to `Task.map()`; defaults to `False`
             - validate (bool, optional): Whether or not to check the validity of
                 the flow (e.g., presence of cycles and illegal keys). Defaults to the value
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
@@ -476,8 +476,8 @@ def set_dependencies(
             - upstream_tasks ([object], optional): A list of upstream tasks for this task
             - downstream_tasks ([object], optional): A list of downtream tasks for this task
             - keyword_tasks ({str, object}}, optional): The results of these tasks will be provided
-            to the task under the specified keyword arguments.
-            - mapped (bool, optional): Whether the results of these tasks should be mapped over
+            to this task under the specified keyword arguments.
+            - mapped (bool, optional): Whether the results of the _upstream_ tasks should be mapped over
                 with the specified keyword arguments
             - validate (bool, optional): Whether or not to check the validity of the flow. If not
                 provided, defaults to the value of `eager_edge_validation` in your Prefect
@@ -504,39 +504,54 @@ def set_dependencies(
             mapped=mapped,
         )
 
-    def set_upstream(self, task: object, flow: "Flow" = None) -> None:
+    def set_upstream(
+        self, task: object, flow: "Flow" = None, key: str = None, mapped: bool = False
+    ) -> None:
         """
         Sets the provided task as an upstream dependency of this task.
 
-        Equivalent to: `self.set_dependencies(upstream_tasks=[task])`
-
         Args:
             - task (object): A task or object that will be converted to a task that will be set
                 as a upstream dependency of this task.
             - flow (Flow, optional): The flow to set dependencies on, defaults to the current
                 flow in context if no flow is specified
+            - key (str, optional): The key to be set for the new edge; the result of the upstream task
+                will be passed to this task's `run()` method under this keyword argument.
+            - mapped (bool, optional): Whether this dependency is mapped; defaults to `False`
 
         Raises:
             - ValueError: if no flow is specified and no flow can be found in the current context
         """
-        self.set_dependencies(flow=flow, upstream_tasks=[task])
+        if key is not None:
+            keyword_tasks = {key: task}
+            self.set_dependencies(flow=flow, keyword_tasks=keyword_tasks, mapped=mapped)
+        else:
+            self.set_dependencies(flow=flow, upstream_tasks=[task], mapped=mapped)
 
-    def set_downstream(self, task: object, flow: "Flow" = None) -> None:
+    def set_downstream(
+        self, task: "Task", flow: "Flow" = None, key: str = None, mapped: bool = False
+    ) -> None:
         """
         Sets the provided task as a downstream dependency of this task.
 
-        Equivalent to: `self.set_dependencies(downstream_tasks=[task])`
-
         Args:
-            - task (object): A task or object that will be converted to a task that will be set
-                as a downstream dependency of this task.
+            - task (Task): A task that will be set as a downstream dependency of this task.
             - flow (Flow, optional): The flow to set dependencies on, defaults to the current
                 flow in context if no flow is specified
+            - key (str, optional): The key to be set for the new edge; the result of this task
+                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`
 
         Raises:
             - ValueError: if no flow is specified and no flow can be found in the current context
         """
-        self.set_dependencies(flow=flow, downstream_tasks=[task])
+        if key is not None:
+            keyword_tasks = {key: self}
+            task.set_dependencies(  # type: ignore
+                flow=flow, keyword_tasks=keyword_tasks, mapped=mapped
+            )  # type: ignore
+        else:
+            task.set_dependencies(flow=flow, upstream_tasks=[self], mapped=mapped)
 
     def inputs(self) -> Dict[str, Dict]:
         """


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
@@ -410,6 +410,16 @@ def test_set_downstream_no_flow(self):
             t1.set_downstream(t2)
         assert "No Flow was passed" in str(exc.value)
 
+    @pytest.mark.parametrize(
+        "props", [{"mapped": True}, {"key": "x"}, {"key": "x", "mapped": True}]
+    )
+    def test_set_downstream_with_properties(self, props):
+        with Flow(name="test") as f:
+            t1 = Task()
+            t2 = Task()
+            t1.set_downstream(t2, **props)
+            assert Edge(t1, t2, **props) in f.edges
+
     def test_set_upstream(self):
         f = Flow(name="test")
         t1 = Task()
@@ -432,6 +442,16 @@ def test_set_upstream_no_flow(self):
             t2.set_upstream(t1)
         assert "No Flow was passed" in str(exc.value)
 
+    @pytest.mark.parametrize(
+        "props", [{"mapped": True}, {"key": "x"}, {"key": "x", "mapped": True}]
+    )
+    def test_set_upstream_with_properties(self, props):
+        with Flow(name="test") as f:
+            t1 = Task()
+            t2 = Task()
+            t2.set_upstream(t1, **props)
+            assert Edge(t1, t2, **props) in f.edges
+
 
 class TestSerialization:
     def test_serialization(self):