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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (96) hide show
  1. prefect/__init__.py +3 -0
  2. prefect/_internal/compatibility/migration.py +1 -1
  3. prefect/_internal/concurrency/api.py +52 -52
  4. prefect/_internal/concurrency/calls.py +59 -35
  5. prefect/_internal/concurrency/cancellation.py +34 -18
  6. prefect/_internal/concurrency/event_loop.py +7 -6
  7. prefect/_internal/concurrency/threads.py +41 -33
  8. prefect/_internal/concurrency/waiters.py +28 -21
  9. prefect/_internal/pydantic/v1_schema.py +2 -2
  10. prefect/_internal/pydantic/v2_schema.py +10 -9
  11. prefect/_internal/schemas/bases.py +10 -11
  12. prefect/_internal/schemas/validators.py +2 -1
  13. prefect/_version.py +3 -3
  14. prefect/automations.py +53 -47
  15. prefect/blocks/abstract.py +12 -10
  16. prefect/blocks/core.py +4 -2
  17. prefect/cache_policies.py +11 -11
  18. prefect/client/__init__.py +3 -1
  19. prefect/client/base.py +36 -37
  20. prefect/client/cloud.py +26 -19
  21. prefect/client/collections.py +2 -2
  22. prefect/client/orchestration.py +366 -277
  23. prefect/client/schemas/__init__.py +24 -0
  24. prefect/client/schemas/actions.py +132 -120
  25. prefect/client/schemas/filters.py +5 -0
  26. prefect/client/schemas/objects.py +113 -85
  27. prefect/client/schemas/responses.py +21 -18
  28. prefect/client/schemas/schedules.py +136 -93
  29. prefect/client/subscriptions.py +28 -14
  30. prefect/client/utilities.py +32 -36
  31. prefect/concurrency/asyncio.py +6 -9
  32. prefect/concurrency/services.py +3 -0
  33. prefect/concurrency/sync.py +35 -5
  34. prefect/context.py +39 -31
  35. prefect/deployments/flow_runs.py +3 -5
  36. prefect/docker/__init__.py +1 -1
  37. prefect/events/schemas/events.py +25 -20
  38. prefect/events/utilities.py +1 -2
  39. prefect/filesystems.py +3 -3
  40. prefect/flow_engine.py +755 -138
  41. prefect/flow_runs.py +3 -3
  42. prefect/flows.py +214 -170
  43. prefect/logging/configuration.py +1 -1
  44. prefect/logging/highlighters.py +1 -2
  45. prefect/logging/loggers.py +30 -20
  46. prefect/main.py +17 -24
  47. prefect/runner/runner.py +43 -21
  48. prefect/runner/server.py +30 -32
  49. prefect/runner/submit.py +3 -6
  50. prefect/runner/utils.py +6 -6
  51. prefect/runtime/flow_run.py +7 -0
  52. prefect/settings/constants.py +2 -2
  53. prefect/settings/legacy.py +1 -1
  54. prefect/settings/models/server/events.py +10 -0
  55. prefect/settings/sources.py +9 -2
  56. prefect/task_engine.py +72 -19
  57. prefect/task_runners.py +2 -2
  58. prefect/tasks.py +46 -33
  59. prefect/telemetry/bootstrap.py +15 -2
  60. prefect/telemetry/run_telemetry.py +107 -0
  61. prefect/transactions.py +14 -14
  62. prefect/types/__init__.py +20 -3
  63. prefect/utilities/_engine.py +96 -0
  64. prefect/utilities/annotations.py +25 -18
  65. prefect/utilities/asyncutils.py +126 -140
  66. prefect/utilities/callables.py +87 -78
  67. prefect/utilities/collections.py +278 -117
  68. prefect/utilities/compat.py +13 -21
  69. prefect/utilities/context.py +6 -5
  70. prefect/utilities/dispatch.py +23 -12
  71. prefect/utilities/dockerutils.py +33 -32
  72. prefect/utilities/engine.py +126 -239
  73. prefect/utilities/filesystem.py +18 -15
  74. prefect/utilities/hashing.py +10 -11
  75. prefect/utilities/importtools.py +40 -27
  76. prefect/utilities/math.py +9 -5
  77. prefect/utilities/names.py +3 -3
  78. prefect/utilities/processutils.py +121 -57
  79. prefect/utilities/pydantic.py +41 -36
  80. prefect/utilities/render_swagger.py +22 -12
  81. prefect/utilities/schema_tools/__init__.py +2 -1
  82. prefect/utilities/schema_tools/hydration.py +50 -43
  83. prefect/utilities/schema_tools/validation.py +52 -42
  84. prefect/utilities/services.py +13 -12
  85. prefect/utilities/templating.py +45 -45
  86. prefect/utilities/text.py +2 -1
  87. prefect/utilities/timeout.py +4 -4
  88. prefect/utilities/urls.py +9 -4
  89. prefect/utilities/visualization.py +46 -24
  90. prefect/variables.py +9 -8
  91. prefect/workers/base.py +18 -10
  92. {prefect_client-3.1.4.dist-info → prefect_client-3.1.6.dist-info}/METADATA +5 -5
  93. {prefect_client-3.1.4.dist-info → prefect_client-3.1.6.dist-info}/RECORD +96 -94
  94. {prefect_client-3.1.4.dist-info → prefect_client-3.1.6.dist-info}/WHEEL +1 -1
  95. {prefect_client-3.1.4.dist-info → prefect_client-3.1.6.dist-info}/LICENSE +0 -0
  96. {prefect_client-3.1.4.dist-info → prefect_client-3.1.6.dist-info}/top_level.txt +0 -0
@@ -5,31 +5,32 @@ Utilities for working with clients.
5
5
  # This module must not import from `prefect.client` when it is imported to avoid
6
6
  # circular imports for decorators such as `inject_client` which are widely used.
7
7
 
8
+ from collections.abc import Awaitable, Coroutine
8
9
  from functools import wraps
9
- from typing import (
10
- TYPE_CHECKING,
11
- Any,
12
- Awaitable,
13
- Callable,
14
- Coroutine,
15
- Optional,
16
- Tuple,
17
- TypeVar,
18
- cast,
19
- )
20
-
21
- from typing_extensions import Concatenate, ParamSpec
10
+ from typing import TYPE_CHECKING, Any, Callable, Optional, Union
11
+
12
+ from typing_extensions import Concatenate, ParamSpec, TypeGuard, TypeVar
22
13
 
23
14
  if TYPE_CHECKING:
24
- from prefect.client.orchestration import PrefectClient
15
+ from prefect.client.orchestration import PrefectClient, SyncPrefectClient
25
16
 
26
17
  P = ParamSpec("P")
27
- R = TypeVar("R")
18
+ R = TypeVar("R", infer_variance=True)
19
+
20
+
21
+ def _current_async_client(
22
+ client: Union["PrefectClient", "SyncPrefectClient"],
23
+ ) -> TypeGuard["PrefectClient"]:
24
+ """Determine if the client is a PrefectClient instance attached to the current loop"""
25
+ from prefect._internal.concurrency.event_loop import get_running_loop
26
+
27
+ # Only a PrefectClient will have a _loop attribute that is the current loop
28
+ return getattr(client, "_loop", None) == get_running_loop()
28
29
 
29
30
 
30
31
  def get_or_create_client(
31
32
  client: Optional["PrefectClient"] = None,
32
- ) -> Tuple["PrefectClient", bool]:
33
+ ) -> tuple["PrefectClient", bool]:
33
34
  """
34
35
  Returns provided client, infers a client from context if available, or creates a new client.
35
36
 
@@ -41,29 +42,22 @@ def get_or_create_client(
41
42
  """
42
43
  if client is not None:
43
44
  return client, True
44
- from prefect._internal.concurrency.event_loop import get_running_loop
45
+
45
46
  from prefect.context import AsyncClientContext, FlowRunContext, TaskRunContext
46
47
 
47
48
  async_client_context = AsyncClientContext.get()
48
49
  flow_run_context = FlowRunContext.get()
49
50
  task_run_context = TaskRunContext.get()
50
51
 
51
- if async_client_context and async_client_context.client._loop == get_running_loop():
52
- return async_client_context.client, True
53
- elif (
54
- flow_run_context
55
- and getattr(flow_run_context.client, "_loop", None) == get_running_loop()
56
- ):
57
- return flow_run_context.client, True
58
- elif (
59
- task_run_context
60
- and getattr(task_run_context.client, "_loop", None) == get_running_loop()
61
- ):
62
- return task_run_context.client, True
63
- else:
64
- from prefect.client.orchestration import get_client as get_httpx_client
52
+ for context in (async_client_context, flow_run_context, task_run_context):
53
+ if context is None:
54
+ continue
55
+ if _current_async_client(context_client := context.client):
56
+ return context_client, True
57
+
58
+ from prefect.client.orchestration import get_client as get_httpx_client
65
59
 
66
- return get_httpx_client(), False
60
+ return get_httpx_client(), False
67
61
 
68
62
 
69
63
  def client_injector(
@@ -90,16 +84,18 @@ def inject_client(
90
84
 
91
85
  @wraps(fn)
92
86
  async def with_injected_client(*args: P.args, **kwargs: P.kwargs) -> R:
93
- client = cast(Optional["PrefectClient"], kwargs.pop("client", None))
94
- client, inferred = get_or_create_client(client)
87
+ given = kwargs.pop("client", None)
88
+ if TYPE_CHECKING:
89
+ assert given is None or isinstance(given, PrefectClient)
90
+ client, inferred = get_or_create_client(given)
95
91
  if not inferred:
96
92
  context = client
97
93
  else:
98
94
  from prefect.utilities.asyncutils import asyncnullcontext
99
95
 
100
- context = asyncnullcontext()
96
+ context = asyncnullcontext(client)
101
97
  async with context as new_client:
102
- kwargs.setdefault("client", new_client or client)
98
+ kwargs |= {"client": new_client}
103
99
  return await fn(*args, **kwargs)
104
100
 
105
101
  return with_injected_client
@@ -17,7 +17,6 @@ except ImportError:
17
17
  from prefect.client.orchestration import get_client
18
18
  from prefect.client.schemas.responses import MinimalConcurrencyLimitResponse
19
19
  from prefect.logging.loggers import get_run_logger
20
- from prefect.utilities.asyncutils import sync_compatible
21
20
 
22
21
  from .context import ConcurrencyContext
23
22
  from .events import (
@@ -79,7 +78,7 @@ async def concurrency(
79
78
 
80
79
  names = names if isinstance(names, list) else [names]
81
80
 
82
- limits = await _acquire_concurrency_slots(
81
+ limits = await _aacquire_concurrency_slots(
83
82
  names,
84
83
  occupy,
85
84
  timeout_seconds=timeout_seconds,
@@ -95,7 +94,7 @@ async def concurrency(
95
94
  finally:
96
95
  occupancy_period = cast(Interval, (pendulum.now("UTC") - acquisition_time))
97
96
  try:
98
- await _release_concurrency_slots(
97
+ await _arelease_concurrency_slots(
99
98
  names, occupy, occupancy_period.total_seconds()
100
99
  )
101
100
  except anyio.get_cancelled_exc_class():
@@ -138,7 +137,7 @@ async def rate_limit(
138
137
 
139
138
  names = names if isinstance(names, list) else [names]
140
139
 
141
- limits = await _acquire_concurrency_slots(
140
+ limits = await _aacquire_concurrency_slots(
142
141
  names,
143
142
  occupy,
144
143
  mode="rate_limit",
@@ -149,7 +148,6 @@ async def rate_limit(
149
148
  _emit_concurrency_acquisition_events(limits, occupy)
150
149
 
151
150
 
152
- @sync_compatible
153
151
  @deprecated_parameter(
154
152
  name="create_if_missing",
155
153
  start_date="Sep 2024",
@@ -157,10 +155,10 @@ async def rate_limit(
157
155
  when=lambda x: x is not None,
158
156
  help="Limits must be explicitly created before acquiring concurrency slots; see `strict` if you want to enforce this behavior.",
159
157
  )
160
- async def _acquire_concurrency_slots(
158
+ async def _aacquire_concurrency_slots(
161
159
  names: List[str],
162
160
  slots: int,
163
- mode: Union[Literal["concurrency"], Literal["rate_limit"]] = "concurrency",
161
+ mode: Literal["concurrency", "rate_limit"] = "concurrency",
164
162
  timeout_seconds: Optional[float] = None,
165
163
  create_if_missing: Optional[bool] = None,
166
164
  max_retries: Optional[int] = None,
@@ -199,8 +197,7 @@ async def _acquire_concurrency_slots(
199
197
  return retval
200
198
 
201
199
 
202
- @sync_compatible
203
- async def _release_concurrency_slots(
200
+ async def _arelease_concurrency_slots(
204
201
  names: List[str], slots: int, occupancy_seconds: float
205
202
  ) -> List[MinimalConcurrencyLimitResponse]:
206
203
  async with get_client() as client:
@@ -83,6 +83,9 @@ class ConcurrencySlotAcquisitionService(QueueService):
83
83
  if max_retries is not None and max_retries <= 0:
84
84
  raise exc
85
85
  retry_after = float(exc.response.headers["Retry-After"])
86
+ logger.debug(
87
+ f"Unable to acquire concurrency slot. Retrying in {retry_after} second(s)."
88
+ )
86
89
  await asyncio.sleep(retry_after)
87
90
  if max_retries is not None:
88
91
  max_retries -= 1
@@ -9,6 +9,9 @@ from typing import (
9
9
  )
10
10
 
11
11
  import pendulum
12
+ from typing_extensions import Literal
13
+
14
+ from prefect.utilities.asyncutils import run_coro_as_sync
12
15
 
13
16
  try:
14
17
  from pendulum import Interval
@@ -19,8 +22,8 @@ except ImportError:
19
22
  from prefect.client.schemas.responses import MinimalConcurrencyLimitResponse
20
23
 
21
24
  from .asyncio import (
22
- _acquire_concurrency_slots,
23
- _release_concurrency_slots,
25
+ _aacquire_concurrency_slots,
26
+ _arelease_concurrency_slots,
24
27
  )
25
28
  from .events import (
26
29
  _emit_concurrency_acquisition_events,
@@ -30,6 +33,36 @@ from .events import (
30
33
  T = TypeVar("T")
31
34
 
32
35
 
36
+ def _release_concurrency_slots(
37
+ names: List[str], slots: int, occupancy_seconds: float
38
+ ) -> List[MinimalConcurrencyLimitResponse]:
39
+ result = run_coro_as_sync(
40
+ _arelease_concurrency_slots(names, slots, occupancy_seconds)
41
+ )
42
+ if result is None:
43
+ raise RuntimeError("Failed to release concurrency slots")
44
+ return result
45
+
46
+
47
+ def _acquire_concurrency_slots(
48
+ names: List[str],
49
+ slots: int,
50
+ mode: Literal["concurrency", "rate_limit"] = "concurrency",
51
+ timeout_seconds: Optional[float] = None,
52
+ create_if_missing: Optional[bool] = None,
53
+ max_retries: Optional[int] = None,
54
+ strict: bool = False,
55
+ ) -> List[MinimalConcurrencyLimitResponse]:
56
+ result = run_coro_as_sync(
57
+ _aacquire_concurrency_slots(
58
+ names, slots, mode, timeout_seconds, create_if_missing, max_retries, strict
59
+ )
60
+ )
61
+ if result is None:
62
+ raise RuntimeError("Failed to acquire concurrency slots")
63
+ return result
64
+
65
+
33
66
  @contextmanager
34
67
  def concurrency(
35
68
  names: Union[str, List[str]],
@@ -81,7 +114,6 @@ def concurrency(
81
114
  create_if_missing=create_if_missing,
82
115
  strict=strict,
83
116
  max_retries=max_retries,
84
- _sync=True,
85
117
  )
86
118
  acquisition_time = pendulum.now("UTC")
87
119
  emitted_events = _emit_concurrency_acquisition_events(limits, occupy)
@@ -94,7 +126,6 @@ def concurrency(
94
126
  names,
95
127
  occupy,
96
128
  occupancy_period.total_seconds(),
97
- _sync=True,
98
129
  )
99
130
  _emit_concurrency_release_events(limits, occupy, emitted_events)
100
131
 
@@ -134,6 +165,5 @@ def rate_limit(
134
165
  timeout_seconds=timeout_seconds,
135
166
  create_if_missing=create_if_missing,
136
167
  strict=strict,
137
- _sync=True,
138
168
  )
139
169
  _emit_concurrency_acquisition_events(limits, occupy)
prefect/context.py CHANGED
@@ -25,7 +25,6 @@ from typing import (
25
25
  Union,
26
26
  )
27
27
 
28
- import pendulum
29
28
  from pydantic import BaseModel, ConfigDict, Field, PrivateAttr
30
29
  from pydantic_extra_types.pendulum_dt import DateTime
31
30
  from typing_extensions import Self
@@ -44,12 +43,16 @@ from prefect.results import (
44
43
  get_default_persist_setting_for_tasks,
45
44
  )
46
45
  from prefect.settings import Profile, Settings
47
- from prefect.settings.legacy import _get_settings_fields
46
+ from prefect.settings.legacy import (
47
+ _get_settings_fields, # type: ignore[reportPrivateUsage]
48
+ )
48
49
  from prefect.states import State
49
50
  from prefect.task_runners import TaskRunner
50
51
  from prefect.utilities.services import start_client_metrics_server
51
52
 
52
53
  T = TypeVar("T")
54
+ P = TypeVar("P")
55
+ R = TypeVar("R")
53
56
 
54
57
  if TYPE_CHECKING:
55
58
  from prefect.flows import Flow
@@ -121,8 +124,8 @@ class ContextModel(BaseModel):
121
124
  """
122
125
 
123
126
  # The context variable for storing data must be defined by the child class
124
- __var__: ContextVar
125
- _token: Optional[Token] = PrivateAttr(None)
127
+ __var__: ContextVar[Self]
128
+ _token: Optional[Token[Self]] = PrivateAttr(None)
126
129
  model_config = ConfigDict(
127
130
  arbitrary_types_allowed=True,
128
131
  extra="forbid",
@@ -150,7 +153,7 @@ class ContextModel(BaseModel):
150
153
  return cls.__var__.get(None)
151
154
 
152
155
  def model_copy(
153
- self: Self, *, update: Optional[Dict[str, Any]] = None, deep: bool = False
156
+ self: Self, *, update: Optional[Mapping[str, Any]] = None, deep: bool = False
154
157
  ):
155
158
  """
156
159
  Duplicate the context model, optionally choosing which fields to include, exclude, or change.
@@ -199,14 +202,14 @@ class SyncClientContext(ContextModel):
199
202
  assert c1 is ctx.client
200
203
  """
201
204
 
202
- __var__ = ContextVar("sync-client-context")
205
+ __var__: ContextVar[Self] = ContextVar("sync-client-context")
203
206
  client: SyncPrefectClient
204
207
  _httpx_settings: Optional[dict[str, Any]] = PrivateAttr(None)
205
208
  _context_stack: int = PrivateAttr(0)
206
209
 
207
210
  def __init__(self, httpx_settings: Optional[dict[str, Any]] = None):
208
211
  super().__init__(
209
- client=get_client(sync_client=True, httpx_settings=httpx_settings),
212
+ client=get_client(sync_client=True, httpx_settings=httpx_settings), # type: ignore[reportCallIssue]
210
213
  )
211
214
  self._httpx_settings = httpx_settings
212
215
  self._context_stack = 0
@@ -220,11 +223,11 @@ class SyncClientContext(ContextModel):
220
223
  else:
221
224
  return self
222
225
 
223
- def __exit__(self, *exc_info):
226
+ def __exit__(self, *exc_info: Any):
224
227
  self._context_stack -= 1
225
228
  if self._context_stack == 0:
226
- self.client.__exit__(*exc_info)
227
- return super().__exit__(*exc_info)
229
+ self.client.__exit__(*exc_info) # type: ignore[reportUnknownMemberType]
230
+ return super().__exit__(*exc_info) # type: ignore[reportUnknownMemberType]
228
231
 
229
232
  @classmethod
230
233
  @contextmanager
@@ -264,12 +267,12 @@ class AsyncClientContext(ContextModel):
264
267
 
265
268
  def __init__(self, httpx_settings: Optional[dict[str, Any]] = None):
266
269
  super().__init__(
267
- client=get_client(sync_client=False, httpx_settings=httpx_settings),
270
+ client=get_client(sync_client=False, httpx_settings=httpx_settings), # type: ignore[reportCallIssue]
268
271
  )
269
272
  self._httpx_settings = httpx_settings
270
273
  self._context_stack = 0
271
274
 
272
- async def __aenter__(self):
275
+ async def __aenter__(self: Self) -> Self:
273
276
  self._context_stack += 1
274
277
  if self._context_stack == 1:
275
278
  await self.client.__aenter__()
@@ -278,11 +281,11 @@ class AsyncClientContext(ContextModel):
278
281
  else:
279
282
  return self
280
283
 
281
- async def __aexit__(self, *exc_info):
284
+ async def __aexit__(self: Self, *exc_info: Any) -> None:
282
285
  self._context_stack -= 1
283
286
  if self._context_stack == 0:
284
- await self.client.__aexit__(*exc_info)
285
- return super().__exit__(*exc_info)
287
+ await self.client.__aexit__(*exc_info) # type: ignore[reportUnknownMemberType]
288
+ return super().__exit__(*exc_info) # type: ignore[reportUnknownMemberType]
286
289
 
287
290
  @classmethod
288
291
  @asynccontextmanager
@@ -305,19 +308,20 @@ class RunContext(ContextModel):
305
308
  client: The Prefect client instance being used for API communication
306
309
  """
307
310
 
308
- def __init__(self, *args, **kwargs):
311
+ def __init__(self, *args: Any, **kwargs: Any):
309
312
  super().__init__(*args, **kwargs)
310
313
 
311
314
  start_client_metrics_server()
312
315
 
313
- start_time: DateTime = Field(default_factory=lambda: pendulum.now("UTC"))
316
+ start_time: DateTime = Field(default_factory=lambda: DateTime.now("UTC"))
314
317
  input_keyset: Optional[Dict[str, Dict[str, str]]] = None
315
318
  client: Union[PrefectClient, SyncPrefectClient]
316
319
 
317
- def serialize(self):
320
+ def serialize(self: Self, include_secrets: bool = True) -> Dict[str, Any]:
318
321
  return self.model_dump(
319
322
  include={"start_time", "input_keyset"},
320
323
  exclude_unset=True,
324
+ context={"include_secrets": include_secrets},
321
325
  )
322
326
 
323
327
 
@@ -336,9 +340,9 @@ class EngineContext(RunContext):
336
340
  flow_run_states: A list of states for flow runs created within this flow run
337
341
  """
338
342
 
339
- flow: Optional["Flow"] = None
343
+ flow: Optional["Flow[Any, Any]"] = None
340
344
  flow_run: Optional[FlowRun] = None
341
- task_runner: TaskRunner
345
+ task_runner: TaskRunner[Any]
342
346
  log_prints: bool = False
343
347
  parameters: Optional[Dict[str, Any]] = None
344
348
 
@@ -351,21 +355,21 @@ class EngineContext(RunContext):
351
355
  persist_result: bool = Field(default_factory=get_default_persist_setting)
352
356
 
353
357
  # Counter for task calls allowing unique
354
- task_run_dynamic_keys: Dict[str, int] = Field(default_factory=dict)
358
+ task_run_dynamic_keys: Dict[str, Union[str, int]] = Field(default_factory=dict)
355
359
 
356
360
  # Counter for flow pauses
357
361
  observed_flow_pauses: Dict[str, int] = Field(default_factory=dict)
358
362
 
359
363
  # Tracking for result from task runs in this flow run for dependency tracking
360
364
  # Holds the ID of the object returned by the task run and task run state
361
- task_run_results: Mapping[int, State] = Field(default_factory=dict)
365
+ task_run_results: dict[int, State] = Field(default_factory=dict)
362
366
 
363
367
  # Events worker to emit events
364
368
  events: Optional[EventsWorker] = None
365
369
 
366
- __var__: ContextVar = ContextVar("flow_run")
370
+ __var__: ContextVar[Self] = ContextVar("flow_run")
367
371
 
368
- def serialize(self):
372
+ def serialize(self: Self, include_secrets: bool = True) -> Dict[str, Any]:
369
373
  return self.model_dump(
370
374
  include={
371
375
  "flow_run",
@@ -378,6 +382,8 @@ class EngineContext(RunContext):
378
382
  "persist_result",
379
383
  },
380
384
  exclude_unset=True,
385
+ serialize_as_any=True,
386
+ context={"include_secrets": include_secrets},
381
387
  )
382
388
 
383
389
 
@@ -394,7 +400,7 @@ class TaskRunContext(RunContext):
394
400
  task_run: The API metadata for this task run
395
401
  """
396
402
 
397
- task: "Task"
403
+ task: "Task[Any, Any]"
398
404
  task_run: TaskRun
399
405
  log_prints: bool = False
400
406
  parameters: Dict[str, Any]
@@ -405,7 +411,7 @@ class TaskRunContext(RunContext):
405
411
 
406
412
  __var__ = ContextVar("task_run")
407
413
 
408
- def serialize(self):
414
+ def serialize(self: Self, include_secrets: bool = True) -> Dict[str, Any]:
409
415
  return self.model_dump(
410
416
  include={
411
417
  "task_run",
@@ -418,6 +424,8 @@ class TaskRunContext(RunContext):
418
424
  "persist_result",
419
425
  },
420
426
  exclude_unset=True,
427
+ serialize_as_any=True,
428
+ context={"include_secrets": include_secrets},
421
429
  )
422
430
 
423
431
 
@@ -436,7 +444,7 @@ class TagsContext(ContextModel):
436
444
  # Return an empty `TagsContext` instead of `None` if no context exists
437
445
  return cls.__var__.get(TagsContext())
438
446
 
439
- __var__: ContextVar = ContextVar("tags")
447
+ __var__: ContextVar[Self] = ContextVar("tags")
440
448
 
441
449
 
442
450
  class SettingsContext(ContextModel):
@@ -453,9 +461,9 @@ class SettingsContext(ContextModel):
453
461
  profile: Profile
454
462
  settings: Settings
455
463
 
456
- __var__: ContextVar = ContextVar("settings")
464
+ __var__: ContextVar[Self] = ContextVar("settings")
457
465
 
458
- def __hash__(self) -> int:
466
+ def __hash__(self: Self) -> int:
459
467
  return hash(self.settings)
460
468
 
461
469
  @classmethod
@@ -562,7 +570,7 @@ def tags(*new_tags: str) -> Generator[Set[str], None, None]:
562
570
 
563
571
  @contextmanager
564
572
  def use_profile(
565
- profile: Union[Profile, str],
573
+ profile: Union[Profile, str, Any],
566
574
  override_environment_variables: bool = False,
567
575
  include_current_context: bool = True,
568
576
  ):
@@ -662,7 +670,7 @@ def root_settings_context():
662
670
  # an override in the `SettingsContext.get` method.
663
671
 
664
672
 
665
- GLOBAL_SETTINGS_CONTEXT: SettingsContext = root_settings_context()
673
+ GLOBAL_SETTINGS_CONTEXT: SettingsContext = root_settings_context() # type: ignore[reportConstantRedefinition]
666
674
 
667
675
 
668
676
  # 2024-07-02: This surfaces an actionable error message for removed objects
@@ -113,10 +113,8 @@ async def run_deployment(
113
113
  task_run_ctx = TaskRunContext.get()
114
114
  if as_subflow and (flow_run_ctx or task_run_ctx):
115
115
  # TODO: this logic can likely be simplified by using `Task.create_run`
116
- from prefect.utilities.engine import (
117
- _dynamic_key_for_task_run,
118
- collect_task_run_inputs,
119
- )
116
+ from prefect.utilities._engine import dynamic_key_for_task_run
117
+ from prefect.utilities.engine import collect_task_run_inputs
120
118
 
121
119
  # This was called from a flow. Link the flow run as a subflow.
122
120
  task_inputs = {
@@ -143,7 +141,7 @@ async def run_deployment(
143
141
  else task_run_ctx.task_run.flow_run_id
144
142
  )
145
143
  dynamic_key = (
146
- _dynamic_key_for_task_run(flow_run_ctx, dummy_task)
144
+ dynamic_key_for_task_run(flow_run_ctx, dummy_task)
147
145
  if flow_run_ctx
148
146
  else task_run_ctx.task_run.dynamic_key
149
147
  )
@@ -17,4 +17,4 @@ def __getattr__(name: str) -> object:
17
17
  module, attr = _public_api[name]
18
18
  return getattr(import_module(module), attr)
19
19
 
20
- raise ImportError(f"module {__name__!r} has no attribute {name!r}")
20
+ raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
@@ -13,16 +13,20 @@ from typing import (
13
13
  )
14
14
  from uuid import UUID, uuid4
15
15
 
16
- import pendulum
17
- from pydantic import ConfigDict, Field, RootModel, field_validator, model_validator
16
+ from pydantic import (
17
+ AfterValidator,
18
+ ConfigDict,
19
+ Field,
20
+ RootModel,
21
+ model_validator,
22
+ )
18
23
  from pydantic_extra_types.pendulum_dt import DateTime
19
- from typing_extensions import Self
24
+ from typing_extensions import Annotated, Self
20
25
 
21
26
  from prefect._internal.schemas.bases import PrefectBaseModel
22
27
  from prefect.logging import get_logger
23
28
  from prefect.settings import (
24
29
  PREFECT_EVENTS_MAXIMUM_LABELS_PER_RESOURCE,
25
- PREFECT_EVENTS_MAXIMUM_RELATED_RESOURCES,
26
30
  )
27
31
 
28
32
  from .labelling import Labelled
@@ -90,22 +94,34 @@ class RelatedResource(Resource):
90
94
  return self["prefect.resource.role"]
91
95
 
92
96
 
97
+ def _validate_related_resources(value) -> List:
98
+ from prefect.settings import PREFECT_EVENTS_MAXIMUM_RELATED_RESOURCES
99
+
100
+ if len(value) > PREFECT_EVENTS_MAXIMUM_RELATED_RESOURCES.value():
101
+ raise ValueError(
102
+ "The maximum number of related resources "
103
+ f"is {PREFECT_EVENTS_MAXIMUM_RELATED_RESOURCES.value()}"
104
+ )
105
+ return value
106
+
107
+
93
108
  class Event(PrefectBaseModel):
94
109
  """The client-side view of an event that has happened to a Resource"""
95
110
 
96
111
  model_config = ConfigDict(extra="ignore")
97
112
 
98
113
  occurred: DateTime = Field(
99
- default_factory=lambda: pendulum.now("UTC"),
114
+ default_factory=lambda: DateTime.now("UTC"),
100
115
  description="When the event happened from the sender's perspective",
101
116
  )
102
- event: str = Field(
103
- description="The name of the event that happened",
104
- )
117
+ event: str = Field(description="The name of the event that happened")
105
118
  resource: Resource = Field(
106
119
  description="The primary Resource this event concerns",
107
120
  )
108
- related: List[RelatedResource] = Field(
121
+ related: Annotated[
122
+ List[RelatedResource],
123
+ AfterValidator(_validate_related_resources),
124
+ ] = Field(
109
125
  default_factory=list,
110
126
  description="A list of additional Resources involved in this event",
111
127
  )
@@ -144,17 +160,6 @@ class Event(PrefectBaseModel):
144
160
  resources[related.role].append(related)
145
161
  return resources
146
162
 
147
- @field_validator("related")
148
- @classmethod
149
- def enforce_maximum_related_resources(cls, value: List[RelatedResource]):
150
- if len(value) > PREFECT_EVENTS_MAXIMUM_RELATED_RESOURCES.value():
151
- raise ValueError(
152
- "The maximum number of related resources "
153
- f"is {PREFECT_EVENTS_MAXIMUM_RELATED_RESOURCES.value()}"
154
- )
155
-
156
- return value
157
-
158
163
  def find_resource_label(self, label: str) -> Optional[str]:
159
164
  """Finds the value of the given label in this event's resource or one of its
160
165
  related resources. If the label starts with `related:<role>:`, search for the
@@ -3,7 +3,6 @@ from typing import Any, Dict, List, Optional, Union
3
3
  from uuid import UUID
4
4
 
5
5
  import pendulum
6
- from pydantic_extra_types.pendulum_dt import DateTime
7
6
 
8
7
  from .clients import (
9
8
  AssertingEventsClient,
@@ -20,7 +19,7 @@ TIGHT_TIMING = timedelta(minutes=5)
20
19
  def emit_event(
21
20
  event: str,
22
21
  resource: Dict[str, str],
23
- occurred: Optional[DateTime] = None,
22
+ occurred: Optional[pendulum.DateTime] = None,
24
23
  related: Optional[Union[List[Dict[str, str]], List[RelatedResource]]] = None,
25
24
  payload: Optional[Dict[str, Any]] = None,
26
25
  id: Optional[UUID] = None,
prefect/filesystems.py CHANGED
@@ -1,6 +1,7 @@
1
1
  import abc
2
2
  import urllib.parse
3
3
  from pathlib import Path
4
+ from shutil import copytree
4
5
  from typing import Any, Dict, Optional
5
6
 
6
7
  import anyio
@@ -13,7 +14,6 @@ from prefect._internal.schemas.validators import (
13
14
  )
14
15
  from prefect.blocks.core import Block
15
16
  from prefect.utilities.asyncutils import run_sync_in_worker_thread, sync_compatible
16
- from prefect.utilities.compat import copytree
17
17
  from prefect.utilities.filesystem import filter_files
18
18
 
19
19
  from ._internal.compatibility.migration import getattr_migration
@@ -158,7 +158,7 @@ class LocalFileSystem(WritableFileSystem, WritableDeploymentStorage):
158
158
  copytree(from_path, local_path, dirs_exist_ok=True, ignore=ignore_func)
159
159
 
160
160
  async def _get_ignore_func(self, local_path: str, ignore_file: str):
161
- with open(ignore_file, "r") as f:
161
+ with open(ignore_file) as f:
162
162
  ignore_patterns = f.readlines()
163
163
  included_files = filter_files(root=local_path, ignore_patterns=ignore_patterns)
164
164
 
@@ -348,7 +348,7 @@ class RemoteFileSystem(WritableFileSystem, WritableDeploymentStorage):
348
348
 
349
349
  included_files = None
350
350
  if ignore_file:
351
- with open(ignore_file, "r") as f:
351
+ with open(ignore_file) as f:
352
352
  ignore_patterns = f.readlines()
353
353
 
354
354
  included_files = filter_files(