Repository: PrefectHQ/prefect

Issue:
Fix up typing hinting on pause/suspend flow run methods when using `wait_for_input`
When using `wait_for_input` the type hinting says you'll get the base `RunInput` back, which makes type checkers upset when writing flows that use it. We should type hint that `pause_flow_run` and `suspend_flow_run` will return an instance of whatever `RunInput` subclass is being given as `wait_for_input`.


Discussion:

Resolving patch:
diff --git a/src/prefect/engine.py b/src/prefect/engine.py
--- a/src/prefect/engine.py
+++ b/src/prefect/engine.py
@@ -103,6 +103,7 @@
     Type,
     TypeVar,
     Union,
+    overload,
 )
 from uuid import UUID, uuid4
 
@@ -207,6 +208,7 @@
 from prefect.utilities.text import truncated_to
 
 R = TypeVar("R")
+T = TypeVar("T", bound=RunInput)
 EngineReturnType = Literal["future", "state", "result"]
 
 
@@ -946,6 +948,30 @@ async def orchestrate_flow_run(
     return state
 
 
+@overload
+async def pause_flow_run(
+    wait_for_input: None = None,
+    flow_run_id: UUID = None,
+    timeout: int = 300,
+    poll_interval: int = 10,
+    reschedule: bool = False,
+    key: str = None,
+) -> None:
+    ...
+
+
+@overload
+async def pause_flow_run(
+    wait_for_input: Type[T],
+    flow_run_id: UUID = None,
+    timeout: int = 300,
+    poll_interval: int = 10,
+    reschedule: bool = False,
+    key: str = None,
+) -> T:
+    ...
+
+
 @sync_compatible
 @deprecated_parameter(
     "flow_run_id", start_date="Dec 2023", help="Use `suspend_flow_run` instead."
@@ -960,12 +986,12 @@ async def orchestrate_flow_run(
     "wait_for_input", group="flow_run_input", when=lambda y: y is not None
 )
 async def pause_flow_run(
+    wait_for_input: Optional[Type[T]] = None,
     flow_run_id: UUID = None,
     timeout: int = 300,
     poll_interval: int = 10,
     reschedule: bool = False,
     key: str = None,
-    wait_for_input: Optional[Type[RunInput]] = None,
 ):
     """
     Pauses the current flow run by blocking execution until resumed.
@@ -1142,14 +1168,36 @@ async def _out_of_process_pause(
         raise RuntimeError(response.details.reason)
 
 
+@overload
+async def suspend_flow_run(
+    wait_for_input: None = None,
+    flow_run_id: Optional[UUID] = None,
+    timeout: Optional[int] = 300,
+    key: Optional[str] = None,
+    client: PrefectClient = None,
+) -> None:
+    ...
+
+
+@overload
+async def suspend_flow_run(
+    wait_for_input: Type[T],
+    flow_run_id: Optional[UUID] = None,
+    timeout: Optional[int] = 300,
+    key: Optional[str] = None,
+    client: PrefectClient = None,
+) -> T:
+    ...
+
+
 @sync_compatible
 @inject_client
 async def suspend_flow_run(
+    wait_for_input: Optional[Type[T]] = None,
     flow_run_id: Optional[UUID] = None,
     timeout: Optional[int] = 300,
     key: Optional[str] = None,
     client: PrefectClient = None,
-    wait_for_input: Optional[Type[RunInput]] = None,
 ):
     """
     Suspends a flow run by stopping code execution until resumed.
diff --git a/src/prefect/input/run_input.py b/src/prefect/input/run_input.py
--- a/src/prefect/input/run_input.py
+++ b/src/prefect/input/run_input.py
@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Type, Union
+from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Type, TypeVar, Union
 from uuid import UUID
 
 import pydantic
@@ -15,6 +15,7 @@
     from prefect._internal.pydantic.v2_schema import create_v2_schema
 
 
+T = TypeVar("T", bound="RunInput")
 KeysetNames = Union[Literal["response"], Literal["schema"]]
 Keyset = Dict[KeysetNames, str]
 
@@ -92,7 +93,7 @@ async def load(cls, keyset: Keyset, flow_run_id: Optional[UUID] = None):
         return cls(**value)
 
     @classmethod
-    def with_initial_data(cls, **kwargs: Any) -> Type["RunInput"]:
+    def with_initial_data(cls: Type[T], **kwargs: Any) -> Type[T]:
         """
         Create a new `RunInput` subclass with the given initial data as field
         defaults.


Test patch: