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

prefect/utilities/schema_tools/hydration.py

@@ -1,10 +1,12 @@
 import json
-from typing import Any, Callable, Dict, Optional
+from abc import ABC, abstractmethod
+from collections.abc import Callable, Sequence
+from typing import Any, Optional, cast
 
 import jinja2
 from pydantic import BaseModel, Field
 from sqlalchemy.ext.asyncio import AsyncSession
-from typing_extensions import TypeAlias
+from typing_extensions import Self, TypeAlias, TypeIs
 
 from prefect.server.utilities.user_templates import (
     TemplateSecurityError,
@@ -15,14 +17,14 @@ from prefect.types import StrictVariableValue
 
 
 class HydrationContext(BaseModel):
-    workspace_variables: Dict[
+    workspace_variables: dict[
         str,
         StrictVariableValue,
     ] = Field(default_factory=dict)
     render_workspace_variables: bool = Field(default=False)
     raise_on_error: bool = Field(default=False)
     render_jinja: bool = Field(default=False)
-    jinja_context: Dict[str, Any] = Field(default_factory=dict)
+    jinja_context: dict[str, Any] = Field(default_factory=dict)
 
     @classmethod
     async def build(
@@ -31,9 +33,11 @@ class HydrationContext(BaseModel):
         raise_on_error: bool = False,
         render_jinja: bool = False,
         render_workspace_variables: bool = False,
-    ) -> "HydrationContext":
+    ) -> Self:
+        from prefect.server.database.orm_models import Variable
         from prefect.server.models.variables import read_variables
 
+        variables: Sequence[Variable]
         if render_workspace_variables:
             variables = await read_variables(
                 session=session,
@@ -51,14 +55,14 @@ class HydrationContext(BaseModel):
         )
 
 
-Handler: TypeAlias = Callable[[dict, HydrationContext], Any]
+Handler: TypeAlias = Callable[[dict[str, Any], HydrationContext], Any]
 PrefectKind: TypeAlias = Optional[str]
 
-_handlers: Dict[PrefectKind, Handler] = {}
+_handlers: dict[PrefectKind, Handler] = {}
 
 
 class Placeholder:
-    def __eq__(self, other):
+    def __eq__(self, other: Any) -> bool:
         return isinstance(other, type(self))
 
     @property
@@ -70,11 +74,11 @@ class RemoveValue(Placeholder):
     pass
 
 
-def _remove_value(value) -> bool:
+def _remove_value(value: Any) -> TypeIs[RemoveValue]:
     return isinstance(value, RemoveValue)
 
 
-class HydrationError(Placeholder, Exception):
+class HydrationError(Placeholder, Exception, ABC):
     def __init__(self, detail: Optional[str] = None):
         self.detail = detail
 
@@ -83,47 +87,49 @@ class HydrationError(Placeholder, Exception):
         return True
 
     @property
-    def message(self):
+    @abstractmethod
+    def message(self) -> str:
         raise NotImplementedError("Must be implemented by subclass")
 
-    def __eq__(self, other):
+    def __eq__(self, other: Any) -> bool:
         return isinstance(other, type(self)) and self.message == other.message
 
-    def __str__(self):
+    def __str__(self) -> str:
         return self.message
 
 
 class KeyNotFound(HydrationError):
     @property
-    def message(self):
+    def message(self) -> str:
         return f"Missing '{self.key}' key in __prefect object"
 
     @property
+    @abstractmethod
     def key(self) -> str:
         raise NotImplementedError("Must be implemented by subclass")
 
 
 class ValueNotFound(KeyNotFound):
     @property
-    def key(self):
+    def key(self) -> str:
         return "value"
 
 
 class TemplateNotFound(KeyNotFound):
     @property
-    def key(self):
+    def key(self) -> str:
         return "template"
 
 
 class VariableNameNotFound(KeyNotFound):
     @property
-    def key(self):
+    def key(self) -> str:
         return "variable_name"
 
 
 class InvalidJSON(HydrationError):
     @property
-    def message(self):
+    def message(self) -> str:
         message = "Invalid JSON"
         if self.detail:
             message += f": {self.detail}"
@@ -132,7 +138,7 @@ class InvalidJSON(HydrationError):
 
 class InvalidJinja(HydrationError):
     @property
-    def message(self):
+    def message(self) -> str:
         message = "Invalid jinja"
         if self.detail:
             message += f": {self.detail}"
@@ -146,29 +152,29 @@ class WorkspaceVariableNotFound(HydrationError):
         return self.detail
 
     @property
-    def message(self):
+    def message(self) -> str:
         return f"Variable '{self.detail}' not found in workspace."
 
 
 class WorkspaceVariable(Placeholder):
-    def __init__(self, variable_name: str):
+    def __init__(self, variable_name: str) -> None:
         self.variable_name = variable_name
 
-    def __eq__(self, other):
+    def __eq__(self, other: Any) -> bool:
         return (
             isinstance(other, type(self)) and self.variable_name == other.variable_name
         )
 
 
 class ValidJinja(Placeholder):
-    def __init__(self, template: str):
+    def __init__(self, template: str) -> None:
         self.template = template
 
-    def __eq__(self, other):
+    def __eq__(self, other: Any) -> bool:
         return isinstance(other, type(self)) and self.template == other.template
 
 
-def handler(kind: PrefectKind) -> Callable:
+def handler(kind: PrefectKind) -> Callable[[Handler], Handler]:
     def decorator(func: Handler) -> Handler:
         _handlers[kind] = func
         return func
@@ -176,9 +182,9 @@ def handler(kind: PrefectKind) -> Callable:
     return decorator
 
 
-def call_handler(kind: PrefectKind, obj: dict, ctx: HydrationContext) -> Any:
+def call_handler(kind: PrefectKind, obj: dict[str, Any], ctx: HydrationContext) -> Any:
     if kind not in _handlers:
-        return (obj or {}).get("value", None)
+        return obj.get("value", None)
 
     res = _handlers[kind](obj, ctx)
     if ctx.raise_on_error and isinstance(res, HydrationError):
@@ -187,7 +193,7 @@ def call_handler(kind: PrefectKind, obj: dict, ctx: HydrationContext) -> Any:
 
 
 @handler("none")
-def null_handler(obj: dict, ctx: HydrationContext):
+def null_handler(obj: dict[str, Any], ctx: HydrationContext):
     if "value" in obj:
         # null handler is a pass through, so we want to continue to hydrate
         return _hydrate(obj["value"], ctx)
@@ -196,7 +202,7 @@ def null_handler(obj: dict, ctx: HydrationContext):
 
 
 @handler("json")
-def json_handler(obj: dict, ctx: HydrationContext):
+def json_handler(obj: dict[str, Any], ctx: HydrationContext):
     if "value" in obj:
         if isinstance(obj["value"], dict):
             dehydrated_json = _hydrate(obj["value"], ctx)
@@ -222,7 +228,7 @@ def json_handler(obj: dict, ctx: HydrationContext):
 
 
 @handler("jinja")
-def jinja_handler(obj: dict, ctx: HydrationContext):
+def jinja_handler(obj: dict[str, Any], ctx: HydrationContext) -> Any:
     if "template" in obj:
         if isinstance(obj["template"], dict):
             dehydrated_jinja = _hydrate(obj["template"], ctx)
@@ -247,7 +253,7 @@ def jinja_handler(obj: dict, ctx: HydrationContext):
 
 
 @handler("workspace_variable")
-def workspace_variable_handler(obj: dict, ctx: HydrationContext):
+def workspace_variable_handler(obj: dict[str, Any], ctx: HydrationContext) -> Any:
     if "variable_name" in obj:
         if isinstance(obj["variable_name"], dict):
             dehydrated_variable = _hydrate(obj["variable_name"], ctx)
@@ -259,7 +265,7 @@ def workspace_variable_handler(obj: dict, ctx: HydrationContext):
             return dehydrated_variable
 
         if not ctx.render_workspace_variables:
-            return WorkspaceVariable(variable_name=obj["variable_name"])
+            return WorkspaceVariable(variable_name=dehydrated_variable)
 
         if dehydrated_variable in ctx.workspace_variables:
             return ctx.workspace_variables[dehydrated_variable]
@@ -277,35 +283,36 @@ def workspace_variable_handler(obj: dict, ctx: HydrationContext):
         return RemoveValue()
 
 
-def hydrate(obj: dict, ctx: Optional[HydrationContext] = None):
-    res = _hydrate(obj, ctx)
+def hydrate(
+    obj: dict[str, Any], ctx: Optional[HydrationContext] = None
+) -> dict[str, Any]:
+    res: dict[str, Any] = _hydrate(obj, ctx)
 
     if _remove_value(res):
-        return {}
+        res = {}
 
     return res
 
 
-def _hydrate(obj, ctx: Optional[HydrationContext] = None) -> Any:
+def _hydrate(obj: Any, ctx: Optional[HydrationContext] = None) -> Any:
     if ctx is None:
         ctx = HydrationContext()
 
-    prefect_object = isinstance(obj, dict) and "__prefect_kind" in obj
-
-    if prefect_object:
-        prefect_kind = obj.get("__prefect_kind")
-        return call_handler(prefect_kind, obj, ctx)
+    if isinstance(obj, dict) and "__prefect_kind" in obj:
+        obj_dict: dict[str, Any] = obj
+        prefect_kind = obj_dict["__prefect_kind"]
+        return call_handler(prefect_kind, obj_dict, ctx)
     else:
         if isinstance(obj, dict):
             return {
                 key: hydrated_value
-                for key, value in obj.items()
+                for key, value in cast(dict[str, Any], obj).items()
                 if not _remove_value(hydrated_value := _hydrate(value, ctx))
             }
         elif isinstance(obj, list):
             return [
                 hydrated_element
-                for element in obj
+                for element in cast(list[Any], obj)
                 if not _remove_value(hydrated_element := _hydrate(element, ctx))
             ]
         else:

prefect/utilities/schema_tools/validation.py

@@ -1,14 +1,19 @@
 from collections import defaultdict, deque
+from collections.abc import Callable, Iterable, Iterator
 from copy import deepcopy
-from typing import Any, Dict, List
+from typing import TYPE_CHECKING, Any, cast
 
 import jsonschema
 from jsonschema.exceptions import ValidationError as JSONSchemaValidationError
 from jsonschema.validators import Draft202012Validator, create
+from referencing.jsonschema import ObjectSchema, Schema
 
 from prefect.utilities.collections import remove_nested_keys
 from prefect.utilities.schema_tools.hydration import HydrationError, Placeholder
 
+if TYPE_CHECKING:
+    from jsonschema.validators import _Validator  # type: ignore
+
 
 class CircularSchemaRefError(Exception):
     pass
@@ -21,12 +26,16 @@ class ValidationError(Exception):
 PLACEHOLDERS_VALIDATOR_NAME = "_placeholders"
 
 
-def _build_validator():
-    def _applicable_validators(schema):
+def _build_validator() -> type["_Validator"]:
+    def _applicable_validators(schema: Schema) -> Iterable[tuple[str, Any]]:
         # the default implementation returns `schema.items()`
-        return {**schema, PLACEHOLDERS_VALIDATOR_NAME: None}.items()
+        assert not isinstance(schema, bool)
+        schema = {**schema, PLACEHOLDERS_VALIDATOR_NAME: None}
+        return schema.items()
 
-    def _placeholders(validator, _, instance, schema):
+    def _placeholders(
+        _validator: "_Validator", _property: object, instance: Any, _schema: Schema
+    ) -> Iterator[JSONSchemaValidationError]:
         if isinstance(instance, HydrationError):
             yield JSONSchemaValidationError(instance.message)
 
@@ -43,7 +52,9 @@ def _build_validator():
         version="prefect",
         type_checker=Draft202012Validator.TYPE_CHECKER,
         format_checker=Draft202012Validator.FORMAT_CHECKER,
-        id_of=Draft202012Validator.ID_OF,
+        id_of=cast(  # the stub for create() is wrong here; id_of accepts (Schema) -> str | None
+            Callable[[Schema], str], Draft202012Validator.ID_OF
+        ),
         applicable_validators=_applicable_validators,
     )
 
@@ -51,24 +62,23 @@ def _build_validator():
 _VALIDATOR = _build_validator()
 
 
-def is_valid_schema(schema: Dict, preprocess: bool = True):
+def is_valid_schema(schema: ObjectSchema, preprocess: bool = True) -> None:
     if preprocess:
         schema = preprocess_schema(schema)
     try:
-        if schema is not None:
-            _VALIDATOR.check_schema(schema, format_checker=_VALIDATOR.FORMAT_CHECKER)
+        _VALIDATOR.check_schema(schema, format_checker=_VALIDATOR.FORMAT_CHECKER)
     except jsonschema.SchemaError as exc:
         raise ValueError(f"Invalid schema: {exc.message}") from exc
 
 
 def validate(
-    obj: Dict,
-    schema: Dict,
+    obj: dict[str, Any],
+    schema: ObjectSchema,
     raise_on_error: bool = False,
     preprocess: bool = True,
     ignore_required: bool = False,
     allow_none_with_default: bool = False,
-) -> List[JSONSchemaValidationError]:
+) -> list[JSONSchemaValidationError]:
     if preprocess:
         schema = preprocess_schema(schema, allow_none_with_default)
 
@@ -93,32 +103,31 @@ def validate(
     else:
         try:
             validator = _VALIDATOR(schema, format_checker=_VALIDATOR.FORMAT_CHECKER)
-            errors = list(validator.iter_errors(obj))
+            errors = list(validator.iter_errors(obj))  # type: ignore
         except RecursionError:
             raise CircularSchemaRefError
         return errors
 
 
-def is_valid(
-    obj: Dict,
-    schema: Dict,
-) -> bool:
+def is_valid(obj: dict[str, Any], schema: ObjectSchema) -> bool:
     errors = validate(obj, schema)
-    return len(errors) == 0
+    return not errors
 
 
-def prioritize_placeholder_errors(errors):
-    errors_by_path = defaultdict(list)
+def prioritize_placeholder_errors(
+    errors: list[JSONSchemaValidationError],
+) -> list[JSONSchemaValidationError]:
+    errors_by_path: dict[str, list[JSONSchemaValidationError]] = defaultdict(list)
     for error in errors:
         path_str = "->".join(str(p) for p in error.relative_path)
         errors_by_path[path_str].append(error)
 
-    filtered_errors = []
-    for path, grouped_errors in errors_by_path.items():
+    filtered_errors: list[JSONSchemaValidationError] = []
+    for grouped_errors in errors_by_path.values():
         placeholders_errors = [
             error
             for error in grouped_errors
-            if error.validator == PLACEHOLDERS_VALIDATOR_NAME
+            if error.validator == PLACEHOLDERS_VALIDATOR_NAME  # type: ignore  # typing stubs are incomplete
         ]
 
         if placeholders_errors:
@@ -129,8 +138,8 @@ def prioritize_placeholder_errors(errors):
     return filtered_errors
 
 
-def build_error_obj(errors: List[JSONSchemaValidationError]) -> Dict:
-    error_response: Dict[str, Any] = {"errors": []}
+def build_error_obj(errors: list[JSONSchemaValidationError]) -> dict[str, Any]:
+    error_response: dict[str, Any] = {"errors": []}
 
     # If multiple errors are present for the same path and one of them
     # is a placeholder error, we want only want to use the placeholder error.
@@ -145,11 +154,11 @@ def build_error_obj(errors: List[JSONSchemaValidationError]) -> Dict:
 
         # Required errors should be moved one level down to the property
         # they're associated with, so we add an extra level to the path.
-        if error.validator == "required":
-            required_field = error.message.split(" ")[0].strip("'")
+        if error.validator == "required":  # type: ignore
+            required_field = error.message.partition(" ")[0].strip("'")
             path.append(required_field)
 
-        current = error_response["errors"]
+        current: list[Any] = error_response["errors"]
 
         # error at the root, just append the error message
         if not path:
@@ -163,10 +172,10 @@ def build_error_obj(errors: List[JSONSchemaValidationError]) -> Dict:
                 else:
                     for entry in current:
                         if entry.get("index") == part:
-                            current = entry["errors"]
+                            current = cast(list[Any], entry["errors"])
                             break
                     else:
-                        new_entry = {"index": part, "errors": []}
+                        new_entry: dict[str, Any] = {"index": part, "errors": []}
                         current.append(new_entry)
                         current = new_entry["errors"]
             else:
@@ -182,7 +191,7 @@ def build_error_obj(errors: List[JSONSchemaValidationError]) -> Dict:
                         current.append(new_entry)
                         current = new_entry["errors"]
 
-    valid = len(error_response["errors"]) == 0
+    valid = not bool(error_response["errors"])
     error_response["valid"] = valid
 
     return error_response
@@ -190,10 +199,10 @@ def build_error_obj(errors: List[JSONSchemaValidationError]) -> Dict:
 
 def _fix_null_typing(
     key: str,
-    schema: Dict,
-    required_fields: List[str],
+    schema: dict[str, Any],
+    required_fields: list[str],
     allow_none_with_default: bool = False,
-):
+) -> None:
     """
     Pydantic V1 does not generate a valid Draft2020-12 schema for null types.
     """
@@ -207,7 +216,7 @@ def _fix_null_typing(
         del schema["type"]
 
 
-def _fix_tuple_items(schema: Dict):
+def _fix_tuple_items(schema: dict[str, Any]) -> None:
     """
     Pydantic V1 does not generate a valid Draft2020-12 schema for tuples.
     """
@@ -216,15 +225,15 @@ def _fix_tuple_items(schema: Dict):
         and isinstance(schema["items"], list)
         and not schema.get("prefixItems")
     ):
-        schema["prefixItems"] = deepcopy(schema["items"])
+        schema["prefixItems"] = deepcopy(cast(list[Any], schema["items"]))
         del schema["items"]
 
 
 def process_properties(
-    properties: Dict,
-    required_fields: List[str],
+    properties: dict[str, dict[str, Any]],
+    required_fields: list[str],
     allow_none_with_default: bool = False,
-):
+) -> None:
     for key, schema in properties.items():
         _fix_null_typing(key, schema, required_fields, allow_none_with_default)
         _fix_tuple_items(schema)
@@ -235,9 +244,9 @@ def process_properties(
 
 
 def preprocess_schema(
-    schema: Dict,
+    schema: ObjectSchema,
     allow_none_with_default: bool = False,
-):
+) -> ObjectSchema:
     schema = deepcopy(schema)
 
     if "properties" in schema:
@@ -247,7 +256,8 @@ def preprocess_schema(
         )
 
     if "definitions" in schema:  # Also process definitions for reused models
-        for definition in (schema["definitions"] or {}).values():
+        definitions = cast(dict[str, Any], schema["definitions"])
+        for definition in definitions.values():
             if "properties" in definition:
                 required_fields = definition.get("required", [])
                 process_properties(

prefect/utilities/services.py

@@ -1,9 +1,10 @@
-import sys
 import threading
 from collections import deque
+from collections.abc import Coroutine
+from logging import Logger
 from traceback import format_exception
 from types import TracebackType
-from typing import Callable, Coroutine, Deque, Optional, Tuple
+from typing import Any, Callable, Optional
 from wsgiref.simple_server import WSGIServer
 
 import anyio
@@ -14,11 +15,11 @@ from prefect.settings import PREFECT_CLIENT_METRICS_ENABLED, PREFECT_CLIENT_METR
 from prefect.utilities.collections import distinct
 from prefect.utilities.math import clamped_poisson_interval
 
-logger = get_logger("utilities.services.critical_service_loop")
+logger: Logger = get_logger("utilities.services.critical_service_loop")
 
 
 async def critical_service_loop(
-    workload: Callable[..., Coroutine],
+    workload: Callable[..., Coroutine[Any, Any, Any]],
     interval: float,
     memory: int = 10,
     consecutive: int = 3,
@@ -26,7 +27,7 @@ async def critical_service_loop(
     printer: Callable[..., None] = print,
     run_once: bool = False,
     jitter_range: Optional[float] = None,
-):
+) -> None:
     """
     Runs the given `workload` function on the specified `interval`, while being
     forgiving of intermittent issues like temporary HTTP errors.  If more than a certain
@@ -50,8 +51,8 @@ async def critical_service_loop(
             between `interval * (1 - range) < rv < interval * (1 + range)`
     """
 
-    track_record: Deque[bool] = deque([True] * consecutive, maxlen=consecutive)
-    failures: Deque[Tuple[Exception, TracebackType]] = deque(maxlen=memory)
+    track_record: deque[bool] = deque([True] * consecutive, maxlen=consecutive)
+    failures: deque[tuple[Exception, Optional[TracebackType]]] = deque(maxlen=memory)
     backoff_count = 0
 
     while True:
@@ -78,7 +79,7 @@ async def critical_service_loop(
             # or Prefect Cloud is having an outage (which will be covered by the
             # exception clause below)
             track_record.append(False)
-            failures.append((exc, sys.exc_info()[-1]))
+            failures.append((exc, exc.__traceback__))
             logger.debug(
                 f"Run of {workload!r} failed with TransportError", exc_info=exc
             )
@@ -88,7 +89,7 @@ async def critical_service_loop(
                 # likely to be temporary and transient.  Don't quit over these unless
                 # it is prolonged.
                 track_record.append(False)
-                failures.append((exc, sys.exc_info()[-1]))
+                failures.append((exc, exc.__traceback__))
                 logger.debug(
                     f"Run of {workload!r} failed with HTTPStatusError", exc_info=exc
                 )
@@ -155,10 +156,10 @@ async def critical_service_loop(
         await anyio.sleep(sleep)
 
 
-_metrics_server: Optional[Tuple[WSGIServer, threading.Thread]] = None
+_metrics_server: Optional[tuple[WSGIServer, threading.Thread]] = None
 
 
-def start_client_metrics_server():
+def start_client_metrics_server() -> None:
     """Start the process-wide Prometheus metrics server for client metrics (if enabled
     with `PREFECT_CLIENT_METRICS_ENABLED`) on the port `PREFECT_CLIENT_METRICS_PORT`."""
     if not PREFECT_CLIENT_METRICS_ENABLED:
@@ -173,7 +174,7 @@ def start_client_metrics_server():
     _metrics_server = start_http_server(port=PREFECT_CLIENT_METRICS_PORT.value())
 
 
-def stop_client_metrics_server():
+def stop_client_metrics_server() -> None:
     """Start the process-wide Prometheus metrics server for client metrics, if it has
     previously been started"""
     global _metrics_server