Repository: PrefectHQ/prefect

Issue:
task.set_upstream() and task.set_downstream() should have a `flow` argument
`task.set_upstream()` and `task.set_downstream()` should have a `flow` argument (which would be passed to `task.set_dependencies()`) so they can be used outside a flow context. They actually reference a passed "flow" in their docstring error message.


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
@@ -35,7 +35,7 @@ def _validate_run_signature(run: Callable) -> None:
             "to *args."
         )
 
-    reserved_kwargs = ["upstream_tasks", "mapped", "task_args"]
+    reserved_kwargs = ["upstream_tasks", "mapped", "task_args", "flow"]
     violations = [kw for kw in reserved_kwargs if kw in run_sig.args]
     if violations:
         msg = "Tasks cannot have the following argument names: {}.".format(
@@ -504,7 +504,7 @@ def set_dependencies(
             mapped=mapped,
         )
 
-    def set_upstream(self, task: object) -> None:
+    def set_upstream(self, task: object, flow: "Flow" = None) -> None:
         """
         Sets the provided task as an upstream dependency of this task.
 
@@ -513,13 +513,15 @@ def set_upstream(self, task: object) -> None:
         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
 
         Raises:
             - ValueError: if no flow is specified and no flow can be found in the current context
         """
-        self.set_dependencies(upstream_tasks=[task])
+        self.set_dependencies(flow=flow, upstream_tasks=[task])
 
-    def set_downstream(self, task: object) -> None:
+    def set_downstream(self, task: object, flow: "Flow" = None) -> None:
         """
         Sets the provided task as a downstream dependency of this task.
 
@@ -528,11 +530,13 @@ def set_downstream(self, task: object) -> None:
         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.
+            - flow (Flow, optional): The flow to set dependencies on, defaults to the current
+                flow in context if no flow is specified
 
         Raises:
             - ValueError: if no flow is specified and no flow can be found in the current context
         """
-        self.set_dependencies(downstream_tasks=[task])
+        self.set_dependencies(flow=flow, downstream_tasks=[task])
 
     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
@@ -111,6 +111,19 @@ class MappedTasks(Task):
                 def run(self, x, mapped=None):
                     pass
 
+    def test_class_instantiation_rejects_mapped_kwarg_decorator(self):
+        with pytest.raises(ValueError):
+
+            @task
+            def run(x, mapped):
+                pass
+
+        with pytest.raises(ValueError):
+
+            @task
+            def run(x, mapped=None):
+                pass
+
     def test_class_instantiation_rejects_upstream_tasks_kwarg(self):
         with pytest.raises(ValueError):
 
@@ -124,6 +137,45 @@ class UpstreamTasks(Task):
                 def run(self, x, upstream_tasks=None):
                     pass
 
+    def test_class_instantiation_rejects_upstream_tasks_kwarg_decorator(self):
+        with pytest.raises(ValueError):
+
+            @task
+            def run(x, upstream_tasks):
+                pass
+
+        with pytest.raises(ValueError):
+
+            @task
+            def run(x, upstream_tasks=None):
+                pass
+
+    def test_class_instantiation_rejects_flow_kwarg(self):
+        with pytest.raises(ValueError):
+
+            class FlowTasks(Task):
+                def run(self, x, flow):
+                    pass
+
+        with pytest.raises(ValueError):
+
+            class FlowTasks(Task):
+                def run(self, x, flow=None):
+                    pass
+
+    def test_class_instantiation_rejects_flow_kwarg_decorator(self):
+        with pytest.raises(ValueError):
+
+            @task
+            def run(x, flow):
+                pass
+
+        with pytest.raises(ValueError):
+
+            @task
+            def run(x, flow=None):
+                pass
+
     def test_class_instantiation_rejects_task_args_kwarg(self):
         with pytest.raises(ValueError):
 
@@ -137,6 +189,19 @@ class TaskArgs(Task):
                 def run(self, x, task_args=None):
                     pass
 
+    def test_class_instantiation_rejects_task_args_kwarg_decorator(self):
+        with pytest.raises(ValueError):
+
+            @task
+            def run(x, task_args):
+                pass
+
+        with pytest.raises(ValueError):
+
+            @task
+            def run(x, task_args=None):
+                pass
+
     def test_create_task_with_and_without_cache_for(self):
         t1 = Task()
         assert t1.cache_validator is never_use
@@ -322,24 +387,50 @@ def test_task_id_must_be_uuid():
 
 
 class TestDependencies:
-    """
-    Most dependnecy tests are done in test_flow.py.
-    """
-
     def test_set_downstream(self):
+        f = Flow(name="test")
+        t1 = Task()
+        t2 = Task()
+        t1.set_downstream(t2, flow=f)
+        assert Edge(t1, t2) in f.edges
+
+    def test_set_downstream_context(self):
         with Flow(name="test") as f:
             t1 = Task()
             t2 = Task()
             t1.set_downstream(t2)
             assert Edge(t1, t2) in f.edges
 
+    def test_set_downstream_no_flow(self):
+        f = Flow(name="test")
+        t1 = Task()
+        t2 = Task()
+        with pytest.raises(ValueError) as exc:
+            t1.set_downstream(t2)
+        assert "No Flow was passed" in str(exc.value)
+
     def test_set_upstream(self):
+        f = Flow(name="test")
+        t1 = Task()
+        t2 = Task()
+        t2.set_upstream(t1, flow=f)
+        assert Edge(t1, t2) in f.edges
+
+    def test_set_upstream_context(self):
         with Flow(name="test") as f:
             t1 = Task()
             t2 = Task()
             t2.set_upstream(t1)
             assert Edge(t1, t2) in f.edges
 
+    def test_set_upstream_no_flow(self):
+        f = Flow(name="test")
+        t1 = Task()
+        t2 = Task()
+        with pytest.raises(ValueError) as exc:
+            t2.set_upstream(t1)
+        assert "No Flow was passed" in str(exc.value)
+
 
 class TestSerialization:
     def test_serialization(self):