more precise typing

PrefectHQ · Oct 23, 2024 · fc43d42 · fc43d42
1 parent d14cd6e
commit fc43d42
Showing 1 changed file with 12 additions and 4 deletions.
diff --git a/src/prefect/tasks.py b/src/prefect/tasks.py
@@ -310,7 +310,9 @@ def __init__(
             Callable[["TaskRunContext", Dict[str, Any]], Optional[str]]
         ] = None,
         cache_expiration: Optional[datetime.timedelta] = None,
-        task_run_name: Optional[Union[Callable[..., str], str]] = None,
+        task_run_name: Optional[
+            Union[Callable[[], str], Callable[[Dict[str, Any]], str], str, Type[NotSet]]
+        ] = NotSet,
         retries: Optional[int] = None,
         retry_delay_seconds: Optional[
             Union[
@@ -531,7 +533,9 @@ def with_options(
         cache_key_fn: Optional[
             Callable[["TaskRunContext", Dict[str, Any]], Optional[str]]
         ] = None,
-        task_run_name: Optional[Union[Callable[..., str], str, Type[NotSet]]] = NotSet,
+        task_run_name: Optional[
+            Union[Callable[[], str], Callable[[Dict[str, Any]], str], str, Type[NotSet]]
+        ] = NotSet,
         cache_expiration: Optional[datetime.timedelta] = None,
         retries: Union[int, Type[NotSet]] = NotSet,
         retry_delay_seconds: Union[
@@ -1583,7 +1587,9 @@ def task(
         Callable[["TaskRunContext", Dict[str, Any]], Optional[str]]
     ] = None,
     cache_expiration: Optional[datetime.timedelta] = None,
-    task_run_name: Optional[Union[Callable[..., str], str]] = None,
+    task_run_name: Optional[
+        Union[Callable[[], str], Callable[[Dict[str, Any]], str], str, Type[NotSet]]
+    ] = NotSet,
     retries: int = 0,
     retry_delay_seconds: Union[
         float,
@@ -1620,7 +1626,9 @@ def task(
         Callable[["TaskRunContext", Dict[str, Any]], Optional[str]], None
     ] = None,
     cache_expiration: Optional[datetime.timedelta] = None,
-    task_run_name: Optional[Union[Callable[..., str], str]] = None,
+    task_run_name: Optional[
+        Union[Callable[[], str], Callable[[Dict[str, Any]], str], str, Type[NotSet]]
+    ] = NotSet,
     retries: Optional[int] = None,
     retry_delay_seconds: Union[
         float, int, List[float], Callable[[int], List[float]], None