prefect-client 3.1.12__py3-none-any.whl → 3.1.14__py3-none-any.whl

prefect/task_runners.py

@@ -40,6 +40,8 @@ from prefect.utilities.callables import (
 from prefect.utilities.collections import isiterable
 
 if TYPE_CHECKING:
+    import logging
+
     from prefect.tasks import Task
 
 P = ParamSpec("P")
@@ -61,11 +63,11 @@ class TaskRunner(abc.ABC, Generic[F]):
     """
 
     def __init__(self):
-        self.logger = get_logger(f"task_runner.{self.name}")
+        self.logger: "logging.Logger" = get_logger(f"task_runner.{self.name}")
         self._started = False
 
     @property
-    def name(self):
+    def name(self) -> str:
         """The name of this task runner"""
         return type(self).__name__.lower().replace("taskrunner", "")
 
@@ -74,32 +76,42 @@ class TaskRunner(abc.ABC, Generic[F]):
         """Return a new instance of this task runner with the same configuration."""
         ...
 
+    @overload
     @abc.abstractmethod
     def submit(
         self,
-        task: "Task[P, R]",
+        task: "Task[P, Coroutine[Any, Any, R]]",
         parameters: dict[str, Any],
         wait_for: Iterable[PrefectFuture[Any]] | None = None,
         dependencies: dict[str, set[TaskRunInput]] | None = None,
     ) -> F:
-        """
-        Submit a task to the task run engine.
+        ...
 
-        Args:
-            task: The task to submit.
-            parameters: The parameters to use when running the task.
-            wait_for: A list of futures that the task depends on.
+    @overload
+    @abc.abstractmethod
+    def submit(
+        self,
+        task: "Task[Any, R]",
+        parameters: dict[str, Any],
+        wait_for: Iterable[PrefectFuture[Any]] | None = None,
+        dependencies: dict[str, set[TaskRunInput]] | None = None,
+    ) -> F:
+        ...
 
-        Returns:
-            A future object that can be used to wait for the task to complete and
-            retrieve the result.
-        """
+    @abc.abstractmethod
+    def submit(
+        self,
+        task: "Task[P, R]",
+        parameters: dict[str, Any],
+        wait_for: Iterable[PrefectFuture[Any]] | None = None,
+        dependencies: dict[str, set[TaskRunInput]] | None = None,
+    ) -> F:
         ...
 
     def map(
         self,
         task: "Task[P, R]",
-        parameters: dict[str, Any],
+        parameters: dict[str, Any | unmapped[Any] | allow_failure[Any]],
         wait_for: Optional[Iterable[PrefectFuture[R]]] = None,
     ) -> PrefectFutureList[F]:
         """
@@ -205,7 +217,7 @@ class TaskRunner(abc.ABC, Generic[F]):
 
         return PrefectFutureList(futures)
 
-    def __enter__(self):
+    def __enter__(self) -> Self:
         if self._started:
             raise RuntimeError("This task runner is already started")
 
@@ -213,12 +225,12 @@ class TaskRunner(abc.ABC, Generic[F]):
         self._started = True
         return self
 
-    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any):
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
         self.logger.debug("Stopping task runner")
         self._started = False
 
 
-class ThreadPoolTaskRunner(TaskRunner[PrefectConcurrentFuture[Any]]):
+class ThreadPoolTaskRunner(TaskRunner[PrefectConcurrentFuture[R]]):
     def __init__(self, max_workers: Optional[int] = None):
         super().__init__()
         self._executor: Optional[ThreadPoolExecutor] = None
@@ -229,7 +241,7 @@ class ThreadPoolTaskRunner(TaskRunner[PrefectConcurrentFuture[Any]]):
         )
         self._cancel_events: Dict[uuid.UUID, threading.Event] = {}
 
-    def duplicate(self) -> "ThreadPoolTaskRunner":
+    def duplicate(self) -> "ThreadPoolTaskRunner[R]":
         return type(self)(max_workers=self._max_workers)
 
     @overload
@@ -254,7 +266,7 @@ class ThreadPoolTaskRunner(TaskRunner[PrefectConcurrentFuture[Any]]):
 
     def submit(
         self,
-        task: "Task[P, R]",
+        task: "Task[P, R | Coroutine[Any, Any, R]]",
         parameters: dict[str, Any],
         wait_for: Iterable[PrefectFuture[Any]] | None = None,
         dependencies: dict[str, set[TaskRunInput]] | None = None,
@@ -345,7 +357,7 @@ class ThreadPoolTaskRunner(TaskRunner[PrefectConcurrentFuture[Any]]):
     ) -> PrefectFutureList[PrefectConcurrentFuture[R]]:
         return super().map(task, parameters, wait_for)
 
-    def cancel_all(self):
+    def cancel_all(self) -> None:
         for event in self._cancel_events.values():
             event.set()
             self.logger.debug("Set cancel event")
@@ -354,12 +366,12 @@ class ThreadPoolTaskRunner(TaskRunner[PrefectConcurrentFuture[Any]]):
             self._executor.shutdown(cancel_futures=True)
             self._executor = None
 
-    def __enter__(self):
+    def __enter__(self) -> Self:
         super().__enter__()
         self._executor = ThreadPoolExecutor(max_workers=self._max_workers)
         return self
 
-    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any):
+    def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
         self.cancel_all()
         if self._executor is not None:
             self._executor.shutdown(cancel_futures=True)
@@ -380,7 +392,7 @@ class PrefectTaskRunner(TaskRunner[PrefectDistributedFuture[R]]):
     def __init__(self):
         super().__init__()
 
-    def duplicate(self) -> "PrefectTaskRunner":
+    def duplicate(self) -> "PrefectTaskRunner[R]":
         return type(self)()
 
     @overload

prefect/task_runs.py

@@ -1,8 +1,10 @@
+from __future__ import annotations
+
 import asyncio
 import atexit
 import threading
 import uuid
-from typing import Callable, Dict, Optional
+from typing import TYPE_CHECKING, Callable, Dict, Optional
 
 import anyio
 from cachetools import TTLCache
@@ -15,6 +17,9 @@ from prefect.events.clients import get_events_subscriber
 from prefect.events.filters import EventFilter, EventNameFilter
 from prefect.logging.loggers import get_logger
 
+if TYPE_CHECKING:
+    import logging
+
 
 class TaskRunWaiter:
     """
@@ -68,19 +73,19 @@ class TaskRunWaiter:
     _instance_lock = threading.Lock()
 
     def __init__(self):
-        self.logger = get_logger("TaskRunWaiter")
-        self._consumer_task: Optional[asyncio.Task] = None
+        self.logger: "logging.Logger" = get_logger("TaskRunWaiter")
+        self._consumer_task: "asyncio.Task[None] | None" = None
         self._observed_completed_task_runs: TTLCache[uuid.UUID, bool] = TTLCache(
             maxsize=10000, ttl=600
         )
         self._completion_events: Dict[uuid.UUID, asyncio.Event] = {}
-        self._completion_callbacks: Dict[uuid.UUID, Callable] = {}
+        self._completion_callbacks: Dict[uuid.UUID, Callable[[], None]] = {}
         self._loop: Optional[asyncio.AbstractEventLoop] = None
         self._observed_completed_task_runs_lock = threading.Lock()
         self._completion_events_lock = threading.Lock()
         self._started = False
 
-    def start(self):
+    def start(self) -> None:
         """
         Start the TaskRunWaiter service.
         """
@@ -89,10 +94,12 @@ class TaskRunWaiter:
         self.logger.debug("Starting TaskRunWaiter")
         loop_thread = get_global_loop()
 
-        if not asyncio.get_running_loop() == loop_thread._loop:
+        if not asyncio.get_running_loop() == loop_thread.loop:
             raise RuntimeError("TaskRunWaiter must run on the global loop thread.")
 
-        self._loop = loop_thread._loop
+        self._loop = loop_thread.loop
+        if TYPE_CHECKING:
+            assert self._loop is not None
 
         consumer_started = asyncio.Event()
         self._consumer_task = self._loop.create_task(
@@ -141,7 +148,7 @@ class TaskRunWaiter:
                 except Exception as exc:
                     self.logger.error(f"Error processing event: {exc}")
 
-    def stop(self):
+    def stop(self) -> None:
         """
         Stop the TaskRunWaiter service.
         """
@@ -155,7 +162,7 @@ class TaskRunWaiter:
     @classmethod
     async def wait_for_task_run(
         cls, task_run_id: uuid.UUID, timeout: Optional[float] = None
-    ):
+    ) -> None:
         """
         Wait for a task run to finish.
 
@@ -199,7 +206,9 @@ class TaskRunWaiter:
                 instance._completion_events.pop(task_run_id, None)
 
     @classmethod
-    def add_done_callback(cls, task_run_id: uuid.UUID, callback):
+    def add_done_callback(
+        cls, task_run_id: uuid.UUID, callback: Callable[[], None]
+    ) -> None:
         """
         Add a callback to be called when a task run finishes.
 
@@ -219,7 +228,7 @@ class TaskRunWaiter:
             instance._completion_callbacks[task_run_id] = callback
 
     @classmethod
-    def instance(cls):
+    def instance(cls) -> Self:
         """
         Get the singleton instance of TaskRunWaiter.
         """

prefect/task_worker.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import asyncio
 import inspect
 import os
@@ -16,7 +18,7 @@ import pendulum
 import uvicorn
 from exceptiongroup import BaseExceptionGroup  # novermin
 from fastapi import FastAPI
-from typing_extensions import ParamSpec, TypeVar
+from typing_extensions import ParamSpec, Self, TypeVar
 from websockets.exceptions import InvalidStatusCode
 
 from prefect import Task
@@ -42,7 +44,10 @@ from prefect.utilities.processutils import (
 from prefect.utilities.services import start_client_metrics_server
 from prefect.utilities.urls import url_for
 
-logger = get_logger("task_worker")
+if TYPE_CHECKING:
+    import logging
+
+logger: "logging.Logger" = get_logger("task_worker")
 
 P = ParamSpec("P")
 R = TypeVar("R", infer_variance=True)
@@ -85,7 +90,7 @@ class TaskWorker:
     def __init__(
         self,
         *tasks: Task[P, R],
-        limit: Optional[int] = 10,
+        limit: int | None = 10,
     ):
         self.tasks: list["Task[..., Any]"] = []
         for t in tasks:
@@ -100,7 +105,7 @@ class TaskWorker:
             else:
                 self.tasks.append(t.with_options(persist_result=True))
 
-        self.task_keys = set(t.task_key for t in tasks if isinstance(t, Task))  # pyright: ignore[reportUnnecessaryIsInstance]
+        self.task_keys: set[str] = set(t.task_key for t in tasks if isinstance(t, Task))  # pyright: ignore[reportUnnecessaryIsInstance]
 
         self._started_at: Optional[pendulum.DateTime] = None
         self.stopping: bool = False
@@ -154,7 +159,7 @@ class TaskWorker:
     def available_tasks(self) -> Optional[int]:
         return int(self._limiter.available_tokens) if self._limiter else None
 
-    def handle_sigterm(self, signum: int, frame: object):
+    def handle_sigterm(self, signum: int, frame: object) -> None:
         """
         Shuts down the task worker when a SIGTERM is received.
         """
@@ -355,14 +360,14 @@ class TaskWorker:
             )
             await asyncio.wrap_future(future)
 
-    async def execute_task_run(self, task_run: TaskRun):
+    async def execute_task_run(self, task_run: TaskRun) -> None:
         """Execute a task run in the task worker."""
         async with self if not self.started else asyncnullcontext():
             token_acquired = await self._acquire_token(task_run.id)
             if token_acquired:
                 await self._safe_submit_scheduled_task_run(task_run)
 
-    async def __aenter__(self):
+    async def __aenter__(self) -> Self:
         logger.debug("Starting task worker...")
 
         if self._client._closed:  # pyright: ignore[reportPrivateUsage]

prefect/tasks.py

@@ -53,10 +53,7 @@ from prefect.results import (
     ResultStore,
     get_or_create_default_task_scheduling_storage,
 )
-from prefect.settings import (
-    PREFECT_TASK_DEFAULT_RETRIES,
-    PREFECT_TASK_DEFAULT_RETRY_DELAY_SECONDS,
-)
+from prefect.settings.context import get_current_settings
 from prefect.states import Pending, Scheduled, State
 from prefect.utilities.annotations import NotSet
 from prefect.utilities.asyncutils import run_coro_as_sync, sync_compatible
@@ -70,6 +67,8 @@ from prefect.utilities.importtools import to_qualified_name
 from prefect.utilities.urls import url_for
 
 if TYPE_CHECKING:
+    import logging
+
     from prefect.client.orchestration import PrefectClient
     from prefect.context import TaskRunContext
     from prefect.transactions import Transaction
@@ -80,7 +79,7 @@ P = ParamSpec("P")  # The parameters of the task
 
 NUM_CHARS_DYNAMIC_KEY = 8
 
-logger = get_logger("tasks")
+logger: "logging.Logger" = get_logger("tasks")
 
 FutureOrResult: TypeAlias = Union[PrefectFuture[T], T]
 OneOrManyFutureOrResult: TypeAlias = Union[
@@ -383,19 +382,19 @@ class Task(Generic[P, R]):
         if not callable(fn):
             raise TypeError("'fn' must be callable")
 
-        self.description = description or inspect.getdoc(fn)
+        self.description: str | None = description or inspect.getdoc(fn)
         update_wrapper(self, fn)
         self.fn = fn
 
         # the task is considered async if its function is async or an async
         # generator
-        self.isasync = asyncio.iscoroutinefunction(
+        self.isasync: bool = asyncio.iscoroutinefunction(
             self.fn
         ) or inspect.isasyncgenfunction(self.fn)
 
         # the task is considered a generator if its function is a generator or
         # an async generator
-        self.isgenerator = inspect.isgeneratorfunction(
+        self.isgenerator: bool = inspect.isgeneratorfunction(
             self.fn
         ) or inspect.isasyncgenfunction(self.fn)
 
@@ -405,7 +404,7 @@ class Task(Generic[P, R]):
             else:
                 self.name = self.fn.__name__
         else:
-            self.name = name
+            self.name: str = name
 
         if task_run_name is not None:
             if not isinstance(task_run_name, str) and not callable(task_run_name):
@@ -420,9 +419,9 @@ class Task(Generic[P, R]):
 
         raise_for_reserved_arguments(self.fn, ["return_state", "wait_for"])
 
-        self.tags = set(tags if tags else [])
+        self.tags: set[str] = set(tags if tags else [])
 
-        self.task_key = _generate_task_key(self.fn)
+        self.task_key: str = _generate_task_key(self.fn)
 
         if cache_policy is not NotSet and cache_key_fn is not None:
             logger.warning(
@@ -464,26 +463,29 @@ class Task(Generic[P, R]):
             # TODO: handle this situation with double storage
             self.cache_policy = None
         else:
-            self.cache_policy = cache_policy
+            self.cache_policy: Union[CachePolicy, type[NotSet], None] = cache_policy
 
         # TaskRunPolicy settings
         # TODO: We can instantiate a `TaskRunPolicy` and add Pydantic bound checks to
         #       validate that the user passes positive numbers here
 
-        self.retries = (
-            retries if retries is not None else PREFECT_TASK_DEFAULT_RETRIES.value()
+        settings = get_current_settings()
+        self.retries: int = (
+            retries if retries is not None else settings.tasks.default_retries
         )
         if retry_delay_seconds is None:
-            retry_delay_seconds = PREFECT_TASK_DEFAULT_RETRY_DELAY_SECONDS.value()
+            retry_delay_seconds = settings.tasks.default_retry_delay_seconds
 
         if callable(retry_delay_seconds):
-            self.retry_delay_seconds = retry_delay_seconds(retries)
+            self.retry_delay_seconds = retry_delay_seconds(self.retries)
         elif not isinstance(retry_delay_seconds, (list, int, float, type(None))):
             raise TypeError(
                 f"Invalid `retry_delay_seconds` provided; must be an int, float, list or callable. Received type {type(retry_delay_seconds)}"
             )
         else:
-            self.retry_delay_seconds = retry_delay_seconds
+            self.retry_delay_seconds: Union[
+                float, int, list[float], None
+            ] = retry_delay_seconds
 
         if isinstance(self.retry_delay_seconds, list) and (
             len(self.retry_delay_seconds) > 50
@@ -507,11 +509,15 @@ class Task(Generic[P, R]):
         self.result_serializer = result_serializer
         self.result_storage_key = result_storage_key
         self.cache_result_in_memory = cache_result_in_memory
-        self.timeout_seconds = float(timeout_seconds) if timeout_seconds else None
-        self.on_rollback_hooks = on_rollback or []
-        self.on_commit_hooks = on_commit or []
-        self.on_completion_hooks = on_completion or []
-        self.on_failure_hooks = on_failure or []
+        self.timeout_seconds: Union[float, None] = (
+            float(timeout_seconds) if timeout_seconds else None
+        )
+        self.on_rollback_hooks: list[Callable[["Transaction"], None]] = (
+            on_rollback or []
+        )
+        self.on_commit_hooks: list[Callable[["Transaction"], None]] = on_commit or []
+        self.on_completion_hooks: list[StateHookCallable] = on_completion or []
+        self.on_failure_hooks: list[StateHookCallable] = on_failure or []
 
         # retry_condition_fn must be a callable or None. If it is neither, raise a TypeError
         if retry_condition_fn is not None and not (callable(retry_condition_fn)):
@@ -527,7 +533,7 @@ class Task(Generic[P, R]):
     def ismethod(self) -> bool:
         return hasattr(self.fn, "__prefect_self__")
 
-    def __get__(self, instance: Any, owner: Any):
+    def __get__(self, instance: Any, owner: Any) -> "Task[P, R]":
         """
         Implement the descriptor protocol so that the task can be used as an instance method.
         When an instance method is loaded, this method is called with the "self" instance as
@@ -582,7 +588,7 @@ class Task(Generic[P, R]):
             Callable[["Task[..., Any]", TaskRun, State], bool]
         ] = None,
         viz_return_value: Optional[Any] = None,
-    ):
+    ) -> "Task[P, R]":
         """
         Create a new task from the current object, updating provided options.
 
@@ -1636,7 +1642,43 @@ def task(
     refresh_cache: Optional[bool] = None,
     on_completion: Optional[list[StateHookCallable]] = None,
     on_failure: Optional[list[StateHookCallable]] = None,
-    retry_condition_fn: Optional[Callable[[Task[P, Any], TaskRun, State], bool]] = None,
+    retry_condition_fn: Literal[None] = None,
+    viz_return_value: Any = None,
+) -> Callable[[Callable[P, R]], Task[P, R]]:
+    ...
+
+
+# see https://github.com/PrefectHQ/prefect/issues/16380
+@overload
+def task(
+    __fn: Literal[None] = None,
+    *,
+    name: Optional[str] = None,
+    description: Optional[str] = None,
+    tags: Optional[Iterable[str]] = None,
+    version: Optional[str] = None,
+    cache_policy: Union[CachePolicy, type[NotSet]] = NotSet,
+    cache_key_fn: Optional[
+        Callable[["TaskRunContext", dict[str, Any]], Optional[str]]
+    ] = None,
+    cache_expiration: Optional[datetime.timedelta] = None,
+    task_run_name: Optional[TaskRunNameValueOrCallable] = None,
+    retries: int = 0,
+    retry_delay_seconds: Union[
+        float, int, list[float], Callable[[int], list[float]], None
+    ] = None,
+    retry_jitter_factor: Optional[float] = None,
+    persist_result: Optional[bool] = None,
+    result_storage: Optional[ResultStorage] = None,
+    result_storage_key: Optional[str] = None,
+    result_serializer: Optional[ResultSerializer] = None,
+    cache_result_in_memory: bool = True,
+    timeout_seconds: Union[int, float, None] = None,
+    log_prints: Optional[bool] = None,
+    refresh_cache: Optional[bool] = None,
+    on_completion: Optional[list[StateHookCallable]] = None,
+    on_failure: Optional[list[StateHookCallable]] = None,
+    retry_condition_fn: Optional[Callable[[Task[P, R], TaskRun, State], bool]] = None,
     viz_return_value: Any = None,
 ) -> Callable[[Callable[P, R]], Task[P, R]]:
     ...

prefect/telemetry/bootstrap.py

@@ -4,7 +4,10 @@ import prefect.settings
 from prefect.client.base import ServerType, determine_server_type
 from prefect.logging.loggers import get_logger
 
-logger = get_logger(__name__)
+if TYPE_CHECKING:
+    import logging
+
+logger: "logging.Logger" = get_logger(__name__)
 
 if TYPE_CHECKING:
     from opentelemetry.sdk._logs import LoggerProvider

prefect/telemetry/run_telemetry.py

@@ -1,6 +1,8 @@
+from __future__ import annotations
+
 import time
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Optional, Union
+from typing import TYPE_CHECKING, Any, Union
 
 from opentelemetry import propagate, trace
 from opentelemetry.context import Context
@@ -47,14 +49,14 @@ class RunTelemetry:
     _tracer: "Tracer" = field(
         default_factory=lambda: get_tracer("prefect", prefect.__version__)
     )
-    span: Optional[Span] = None
+    span: Span | None = None
 
     async def async_start_span(
         self,
         run: FlowOrTaskRun,
         client: PrefectClient,
-        parameters: Optional[dict[str, Any]] = None,
-    ):
+        parameters: dict[str, Any] | None = None,
+    ) -> Span:
         traceparent, span = self._start_span(run, parameters)
 
         if self._run_type(run) == "flow" and traceparent:
@@ -70,8 +72,8 @@ class RunTelemetry:
         self,
         run: FlowOrTaskRun,
         client: SyncPrefectClient,
-        parameters: Optional[dict[str, Any]] = None,
-    ):
+        parameters: dict[str, Any] | None = None,
+    ) -> Span:
         traceparent, span = self._start_span(run, parameters)
 
         if self._run_type(run) == "flow" and traceparent:
@@ -84,8 +86,8 @@ class RunTelemetry:
     def _start_span(
         self,
         run: FlowOrTaskRun,
-        parameters: Optional[dict[str, Any]] = None,
-    ) -> tuple[Optional[str], Span]:
+        parameters: dict[str, Any] | None = None,
+    ) -> tuple[str | None, Span]:
         """
         Start a span for a run.
         """
@@ -139,8 +141,8 @@ class RunTelemetry:
         return "task" if isinstance(run, TaskRun) else "flow"
 
     def _trace_context_from_labels(
-        self, labels: Optional[KeyValueLabels]
-    ) -> Optional[Context]:
+        self, labels: KeyValueLabels | None
+    ) -> Context | None:
         """Get trace context from run labels if it exists."""
         if not labels or LABELS_TRACEPARENT_KEY not in labels:
             return None
@@ -148,7 +150,7 @@ class RunTelemetry:
         carrier = {TRACEPARENT_KEY: traceparent}
         return propagate.extract(carrier)
 
-    def _traceparent_from_span(self, span: Span) -> Optional[str]:
+    def _traceparent_from_span(self, span: Span) -> str | None:
         carrier: dict[str, Any] = {}
         propagate.inject(carrier, context=trace.set_span_in_context(span))
         return carrier.get(TRACEPARENT_KEY)
@@ -162,7 +164,7 @@ class RunTelemetry:
             self.span.end(time.time_ns())
             self.span = None
 
-    def end_span_on_failure(self, terminal_message: Optional[str] = None) -> None:
+    def end_span_on_failure(self, terminal_message: str | None = None) -> None:
         """
         End a span for a run on failure.
         """
@@ -203,7 +205,7 @@ class RunTelemetry:
             self.span.update_name(name=name)
             self.span.set_attribute("prefect.run.name", name)
 
-    def _parent_run(self) -> Union[FlowOrTaskRun, None]:
+    def _parent_run(self) -> FlowOrTaskRun | None:
         """
         Identify the "parent run" for the current execution context.
 

prefect/transactions.py

@@ -173,7 +173,7 @@ class Transaction(ContextModel):
     def is_active(self) -> bool:
         return self.state == TransactionState.ACTIVE
 
-    def __enter__(self):
+    def __enter__(self) -> Self:
         if self._token is not None:
             raise RuntimeError(
                 "Context already entered. Context enter calls cannot be nested."
@@ -206,7 +206,7 @@ class Transaction(ContextModel):
         self._token = self.__var__.set(self)
         return self
 
-    def __exit__(self, *exc_info: Any):
+    def __exit__(self, *exc_info: Any) -> None:
         exc_type, exc_val, _ = exc_info
         if not self._token:
             raise RuntimeError(
@@ -235,7 +235,7 @@ class Transaction(ContextModel):
 
         self.reset()
 
-    def begin(self):
+    def begin(self) -> None:
         if (
             self.store
             and self.key

prefect/types/__init__.py

@@ -1,10 +1,13 @@
+from __future__ import annotations
+
 from functools import partial
 from typing import Annotated, Any, Dict, List, Optional, Set, TypeVar, Union
-from typing_extensions import Literal, TypeAlias
+from typing_extensions import Literal
 import orjson
 import pydantic
-from pydantic_extra_types.pendulum_dt import DateTime as PydanticDateTime
-from pydantic_extra_types.pendulum_dt import Date as PydanticDate
+
+
+from ._datetime import DateTime, Date
 from pydantic import (
     BeforeValidator,
     Field,
@@ -35,8 +38,6 @@ TimeZone = Annotated[
     ),
 ]
 
-DateTime: TypeAlias = PydanticDateTime
-Date: TypeAlias = PydanticDate
 
 BANNED_CHARACTERS = ["/", "%", "&", ">", "<"]
 
@@ -114,7 +115,7 @@ class SecretDict(pydantic.Secret[Dict[str, Any]]):
 
 
 def validate_set_T_from_delim_string(
-    value: Union[str, T, Set[T], None], type_, delim=None
+    value: Union[str, T, Set[T], None], type_: type[T], delim: str | None = None
 ) -> Set[T]:
     """
     "no-info" before validator useful in scooping env vars
@@ -169,6 +170,8 @@ KeyValueLabelsField = Annotated[
 
 __all__ = [
     "ClientRetryExtraCodes",
+    "Date",
+    "DateTime",
     "LogLevel",
     "KeyValueLabelsField",
     "NonNegativeInteger",