add overloads for run_task (#17448)

PrefectHQ · Mar 11, 2025 · b958b0f · b958b0f
1 parent 1fedfca
commit b958b0f
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 0 deletions.
diff --git a/src/prefect/task_engine.py b/src/prefect/task_engine.py
@@ -24,6 +24,7 @@
     Type,
     TypeVar,
     Union,
+    overload,
 )
 from uuid import UUID
 
@@ -1527,6 +1528,32 @@ async def run_generator_task_async(
         await engine.result()
 
 
+@overload
+def run_task(
+    task: "Task[P, R]",
+    task_run_id: Optional[UUID] = None,
+    task_run: Optional[TaskRun] = None,
+    parameters: Optional[dict[str, Any]] = None,
+    wait_for: Optional["OneOrManyFutureOrResult[Any]"] = None,
+    return_type: Literal["state"] = "state",
+    dependencies: Optional[dict[str, set[TaskRunInput]]] = None,
+    context: Optional[dict[str, Any]] = None,
+) -> State[R]: ...
+
+
+@overload
+def run_task(
+    task: "Task[P, R]",
+    task_run_id: Optional[UUID] = None,
+    task_run: Optional[TaskRun] = None,
+    parameters: Optional[dict[str, Any]] = None,
+    wait_for: Optional["OneOrManyFutureOrResult[Any]"] = None,
+    return_type: Literal["result"] = "result",
+    dependencies: Optional[dict[str, set[TaskRunInput]]] = None,
+    context: Optional[dict[str, Any]] = None,
+) -> R: ...
+
+
 def run_task(
     task: "Task[P, Union[R, Coroutine[Any, Any, R]]]",
     task_run_id: Optional[UUID] = None,

diff --git a/src/prefect/tasks.py b/src/prefect/tasks.py
@@ -1426,6 +1426,8 @@ def map(
         else:
             return futures
 
+    # Background task methods
+
     def apply_async(
         self,
         args: Optional[tuple[Any, ...]] = None,