prefect-client 3.1.5__py3-none-any.whl → 3.1.6__py3-none-any.whl

prefect/transactions.py

@@ -23,7 +23,7 @@ from prefect.exceptions import (
     MissingContextError,
     SerializationError,
 )
-from prefect.logging.loggers import get_logger, get_run_logger
+from prefect.logging.loggers import LoggingAdapter, get_logger, get_run_logger
 from prefect.records import RecordStore
 from prefect.records.base import TransactionRecord
 from prefect.results import (
@@ -32,9 +32,9 @@ from prefect.results import (
     ResultStore,
     get_result_store,
 )
+from prefect.utilities._engine import get_hook_name
 from prefect.utilities.annotations import NotSet
 from prefect.utilities.collections import AutoEnum
-from prefect.utilities.engine import _get_hook_name
 
 
 class IsolationLevel(AutoEnum):
@@ -72,13 +72,13 @@ class Transaction(ContextModel):
         default_factory=list
     )
     overwrite: bool = False
-    logger: Union[logging.Logger, logging.LoggerAdapter] = Field(
+    logger: Union[logging.Logger, LoggingAdapter] = Field(
         default_factory=partial(get_logger, "transactions")
     )
     write_on_commit: bool = True
     _stored_values: Dict[str, Any] = PrivateAttr(default_factory=dict)
     _staged_value: Any = None
-    __var__: ContextVar = ContextVar("transaction")
+    __var__: ContextVar[Self] = ContextVar("transaction")
 
     def set(self, name: str, value: Any) -> None:
         """
@@ -209,7 +209,7 @@ class Transaction(ContextModel):
         self._token = self.__var__.set(self)
         return self
 
-    def __exit__(self, *exc_info):
+    def __exit__(self, *exc_info: Any):
         exc_type, exc_val, _ = exc_info
         if not self._token:
             raise RuntimeError(
@@ -254,7 +254,7 @@ class Transaction(ContextModel):
         ):
             self.state = TransactionState.COMMITTED
 
-    def read(self) -> Union["BaseResult", ResultRecord, None]:
+    def read(self) -> Union["BaseResult[Any]", ResultRecord[Any], None]:
         if self.store and self.key:
             record = self.store.read(key=self.key)
             if isinstance(record, ResultRecord):
@@ -354,8 +354,8 @@ class Transaction(ContextModel):
             self.rollback()
             return False
 
-    def run_hook(self, hook, hook_type: str) -> None:
-        hook_name = _get_hook_name(hook)
+    def run_hook(self, hook: Callable[..., Any], hook_type: str) -> None:
+        hook_name = get_hook_name(hook)
         # Undocumented way to disable logging for a hook. Subject to change.
         should_log = getattr(hook, "log_on_run", True)
 
@@ -379,8 +379,8 @@ class Transaction(ContextModel):
     def stage(
         self,
         value: Any,
-        on_rollback_hooks: Optional[List] = None,
-        on_commit_hooks: Optional[List] = None,
+        on_rollback_hooks: Optional[list[Callable[..., Any]]] = None,
+        on_commit_hooks: Optional[list[Callable[..., Any]]] = None,
     ) -> None:
         """
         Stage a value to be committed later.
@@ -441,7 +441,7 @@ def transaction(
     isolation_level: Optional[IsolationLevel] = None,
     overwrite: bool = False,
     write_on_commit: bool = True,
-    logger: Union[logging.Logger, logging.LoggerAdapter, None] = None,
+    logger: Optional[Union[logging.Logger, LoggingAdapter]] = None,
 ) -> Generator[Transaction, None, None]:
     """
     A context manager for opening and managing a transaction.
@@ -465,9 +465,9 @@ def transaction(
         store = get_result_store()
 
     try:
-        logger = logger or get_run_logger()
+        _logger: Union[logging.Logger, LoggingAdapter] = logger or get_run_logger()
     except MissingContextError:
-        logger = get_logger("transactions")
+        _logger = get_logger("transactions")
 
     with Transaction(
         key=key,
@@ -476,6 +476,6 @@ def transaction(
         isolation_level=isolation_level,
         overwrite=overwrite,
         write_on_commit=write_on_commit,
-        logger=logger,
+        logger=_logger,
     ) as txn:
         yield txn

prefect/types/__init__.py

@@ -96,7 +96,7 @@ def cast_none_to_empty_dict(value: Any) -> dict[str, Any]:
 
 
 KeyValueLabels = Annotated[
-    dict[str, Union[StrictBool, StrictInt, StrictFloat, str]],
+    Dict[str, Union[StrictBool, StrictInt, StrictFloat, str]],
     BeforeValidator(cast_none_to_empty_dict),
 ]
 
@@ -149,9 +149,6 @@ LogLevel = Annotated[
 ]
 
 
-KeyValueLabels: TypeAlias = dict[str, Union[StrictBool, StrictInt, StrictFloat, str]]
-
-
 def convert_none_to_empty_dict(v: Optional[KeyValueLabels]) -> KeyValueLabels:
     return v or {}
 

prefect/utilities/_engine.py

@@ -0,0 +1,96 @@
+"""Internal engine utilities"""
+
+
+from collections.abc import Callable
+from functools import partial
+from typing import TYPE_CHECKING, Any, Union
+from uuid import uuid4
+
+from prefect.context import FlowRunContext
+from prefect.flows import Flow
+from prefect.tasks import Task, TaskRunNameCallbackWithParameters
+
+
+def dynamic_key_for_task_run(
+    context: FlowRunContext, task: "Task[..., Any]", stable: bool = True
+) -> Union[int, str]:
+    if (
+        stable is False or context.detached
+    ):  # this task is running on remote infrastructure
+        return str(uuid4())
+    elif context.flow_run is None:  # this is an autonomous task run
+        context.task_run_dynamic_keys[task.task_key] = getattr(
+            task, "dynamic_key", str(uuid4())
+        )
+
+    elif task.task_key not in context.task_run_dynamic_keys:
+        context.task_run_dynamic_keys[task.task_key] = 0
+    else:
+        dynamic_key = context.task_run_dynamic_keys[task.task_key]
+        if TYPE_CHECKING:
+            assert isinstance(dynamic_key, int)
+        context.task_run_dynamic_keys[task.task_key] = dynamic_key + 1
+
+    return context.task_run_dynamic_keys[task.task_key]
+
+
+def resolve_custom_flow_run_name(
+    flow: "Flow[..., Any]", parameters: dict[str, Any]
+) -> str:
+    if callable(flow.flow_run_name):
+        flow_run_name = flow.flow_run_name()
+        if not TYPE_CHECKING:
+            if not isinstance(flow_run_name, str):
+                raise TypeError(
+                    f"Callable {flow.flow_run_name} for 'flow_run_name' returned type"
+                    f" {type(flow_run_name).__name__} but a string is required."
+                )
+    elif isinstance(flow.flow_run_name, str):
+        flow_run_name = flow.flow_run_name.format(**parameters)
+    else:
+        raise TypeError(
+            "Expected string or callable for 'flow_run_name'; got"
+            f" {type(flow.flow_run_name).__name__} instead."
+        )
+
+    return flow_run_name
+
+
+def resolve_custom_task_run_name(
+    task: "Task[..., Any]", parameters: dict[str, Any]
+) -> str:
+    if callable(task.task_run_name):
+        # If the callable accepts a 'parameters' kwarg, pass the entire parameters dict
+        if TaskRunNameCallbackWithParameters.is_callback_with_parameters(
+            task.task_run_name
+        ):
+            task_run_name = task.task_run_name(parameters=parameters)
+        else:
+            # If it doesn't expect parameters, call it without arguments
+            task_run_name = task.task_run_name()
+
+        if not TYPE_CHECKING:
+            if not isinstance(task_run_name, str):
+                raise TypeError(
+                    f"Callable {task.task_run_name} for 'task_run_name' returned type"
+                    f" {type(task_run_name).__name__} but a string is required."
+                )
+    elif isinstance(task.task_run_name, str):
+        task_run_name = task.task_run_name.format(**parameters)
+    else:
+        raise TypeError(
+            "Expected string or callable for 'task_run_name'; got"
+            f" {type(task.task_run_name).__name__} instead."
+        )
+
+    return task_run_name
+
+
+def get_hook_name(hook: Callable[..., Any]) -> str:
+    return (
+        hook.__name__
+        if hasattr(hook, "__name__")
+        else (
+            hook.func.__name__ if isinstance(hook, partial) else hook.__class__.__name__
+        )
+    )

prefect/utilities/annotations.py

@@ -1,33 +1,40 @@
 import warnings
-from abc import ABC
-from collections import namedtuple
-from typing import Generic, TypeVar
+from operator import itemgetter
+from typing import Any, cast
 
-T = TypeVar("T")
+from typing_extensions import Self, TypeVar
 
+T = TypeVar("T", infer_variance=True)
 
-class BaseAnnotation(
-    namedtuple("BaseAnnotation", field_names="value"), ABC, Generic[T]
-):
+
+class BaseAnnotation(tuple[T]):
     """
     Base class for Prefect annotation types.
 
-    Inherits from `namedtuple` for unpacking support in another tools.
+    Inherits from `tuple` for unpacking support in other tools.
     """
 
+    __slots__ = ()
+
+    def __new__(cls, value: T) -> Self:
+        return super().__new__(cls, (value,))
+
+    # use itemgetter to minimise overhead, just like namedtuple generated code would
+    value: T = cast(T, property(itemgetter(0)))
+
     def unwrap(self) -> T:
-        return self.value
+        return self[0]
 
-    def rewrap(self, value: T) -> "BaseAnnotation[T]":
+    def rewrap(self, value: T) -> Self:
         return type(self)(value)
 
-    def __eq__(self, other: "BaseAnnotation[T]") -> bool:
+    def __eq__(self, other: Any) -> bool:
         if type(self) is not type(other):
             return False
-        return self.unwrap() == other.unwrap()
+        return super().__eq__(other)
 
     def __repr__(self) -> str:
-        return f"{type(self).__name__}({self.value!r})"
+        return f"{type(self).__name__}({self[0]!r})"
 
 
 class unmapped(BaseAnnotation[T]):
@@ -38,9 +45,9 @@ class unmapped(BaseAnnotation[T]):
     operation instead of being split.
     """
 
-    def __getitem__(self, _) -> T:
+    def __getitem__(self, _: object) -> T:  # type: ignore[override]  # pyright: ignore[reportIncompatibleMethodOverride]
         # Internally, this acts as an infinite array where all items are the same value
-        return self.unwrap()
+        return super().__getitem__(0)
 
 
 class allow_failure(BaseAnnotation[T]):
@@ -87,14 +94,14 @@ class quote(BaseAnnotation[T]):
 
 
 # Backwards compatibility stub for `Quote` class
-class Quote(quote):
-    def __init__(self, expr):
+class Quote(quote[T]):
+    def __new__(cls, expr: T) -> Self:
         warnings.warn(
             "Use of `Quote` is deprecated. Use `quote` instead.",
             DeprecationWarning,
             stacklevel=2,
         )
-        super().__init__(expr)
+        return super().__new__(cls, expr)
 
 
 class NotSet: