prefect-client 3.1.10__py3-none-any.whl → 3.1.12__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 (141) hide show
  1. prefect/_experimental/lineage.py +7 -8
  2. prefect/_experimental/sla/__init__.py +0 -0
  3. prefect/_experimental/sla/client.py +66 -0
  4. prefect/_experimental/sla/objects.py +53 -0
  5. prefect/_internal/_logging.py +15 -3
  6. prefect/_internal/compatibility/async_dispatch.py +22 -16
  7. prefect/_internal/compatibility/deprecated.py +42 -18
  8. prefect/_internal/compatibility/migration.py +2 -2
  9. prefect/_internal/concurrency/inspection.py +12 -14
  10. prefect/_internal/concurrency/primitives.py +2 -2
  11. prefect/_internal/concurrency/services.py +154 -80
  12. prefect/_internal/concurrency/waiters.py +13 -9
  13. prefect/_internal/pydantic/annotations/pendulum.py +7 -7
  14. prefect/_internal/pytz.py +4 -3
  15. prefect/_internal/retries.py +10 -5
  16. prefect/_internal/schemas/bases.py +19 -10
  17. prefect/_internal/schemas/validators.py +227 -388
  18. prefect/_version.py +3 -3
  19. prefect/automations.py +236 -30
  20. prefect/blocks/__init__.py +3 -3
  21. prefect/blocks/abstract.py +53 -30
  22. prefect/blocks/core.py +183 -84
  23. prefect/blocks/notifications.py +133 -73
  24. prefect/blocks/redis.py +13 -9
  25. prefect/blocks/system.py +24 -11
  26. prefect/blocks/webhook.py +7 -5
  27. prefect/cache_policies.py +3 -2
  28. prefect/client/orchestration/__init__.py +1957 -0
  29. prefect/client/orchestration/_artifacts/__init__.py +0 -0
  30. prefect/client/orchestration/_artifacts/client.py +239 -0
  31. prefect/client/orchestration/_automations/__init__.py +0 -0
  32. prefect/client/orchestration/_automations/client.py +329 -0
  33. prefect/client/orchestration/_blocks_documents/__init__.py +0 -0
  34. prefect/client/orchestration/_blocks_documents/client.py +334 -0
  35. prefect/client/orchestration/_blocks_schemas/__init__.py +0 -0
  36. prefect/client/orchestration/_blocks_schemas/client.py +200 -0
  37. prefect/client/orchestration/_blocks_types/__init__.py +0 -0
  38. prefect/client/orchestration/_blocks_types/client.py +380 -0
  39. prefect/client/orchestration/_concurrency_limits/__init__.py +0 -0
  40. prefect/client/orchestration/_concurrency_limits/client.py +762 -0
  41. prefect/client/orchestration/_deployments/__init__.py +0 -0
  42. prefect/client/orchestration/_deployments/client.py +1128 -0
  43. prefect/client/orchestration/_flow_runs/__init__.py +0 -0
  44. prefect/client/orchestration/_flow_runs/client.py +903 -0
  45. prefect/client/orchestration/_flows/__init__.py +0 -0
  46. prefect/client/orchestration/_flows/client.py +343 -0
  47. prefect/client/orchestration/_logs/__init__.py +0 -0
  48. prefect/client/orchestration/_logs/client.py +97 -0
  49. prefect/client/orchestration/_variables/__init__.py +0 -0
  50. prefect/client/orchestration/_variables/client.py +157 -0
  51. prefect/client/orchestration/base.py +46 -0
  52. prefect/client/orchestration/routes.py +145 -0
  53. prefect/client/schemas/__init__.py +68 -28
  54. prefect/client/schemas/actions.py +2 -2
  55. prefect/client/schemas/filters.py +5 -0
  56. prefect/client/schemas/objects.py +8 -15
  57. prefect/client/schemas/schedules.py +22 -10
  58. prefect/concurrency/_asyncio.py +87 -0
  59. prefect/concurrency/{events.py → _events.py} +10 -10
  60. prefect/concurrency/asyncio.py +20 -104
  61. prefect/concurrency/context.py +6 -4
  62. prefect/concurrency/services.py +26 -74
  63. prefect/concurrency/sync.py +23 -44
  64. prefect/concurrency/v1/_asyncio.py +63 -0
  65. prefect/concurrency/v1/{events.py → _events.py} +13 -15
  66. prefect/concurrency/v1/asyncio.py +27 -80
  67. prefect/concurrency/v1/context.py +6 -4
  68. prefect/concurrency/v1/services.py +33 -79
  69. prefect/concurrency/v1/sync.py +18 -37
  70. prefect/context.py +66 -45
  71. prefect/deployments/base.py +10 -144
  72. prefect/deployments/flow_runs.py +12 -2
  73. prefect/deployments/runner.py +53 -4
  74. prefect/deployments/steps/pull.py +13 -0
  75. prefect/engine.py +17 -4
  76. prefect/events/clients.py +7 -1
  77. prefect/events/schemas/events.py +3 -2
  78. prefect/filesystems.py +6 -2
  79. prefect/flow_engine.py +101 -85
  80. prefect/flows.py +10 -1
  81. prefect/input/run_input.py +2 -1
  82. prefect/logging/logging.yml +1 -1
  83. prefect/main.py +1 -3
  84. prefect/results.py +2 -307
  85. prefect/runner/runner.py +4 -2
  86. prefect/runner/storage.py +87 -21
  87. prefect/serializers.py +32 -25
  88. prefect/settings/legacy.py +4 -4
  89. prefect/settings/models/api.py +3 -3
  90. prefect/settings/models/cli.py +3 -3
  91. prefect/settings/models/client.py +5 -3
  92. prefect/settings/models/cloud.py +8 -3
  93. prefect/settings/models/deployments.py +3 -3
  94. prefect/settings/models/experiments.py +4 -7
  95. prefect/settings/models/flows.py +3 -3
  96. prefect/settings/models/internal.py +4 -2
  97. prefect/settings/models/logging.py +4 -3
  98. prefect/settings/models/results.py +3 -3
  99. prefect/settings/models/root.py +3 -2
  100. prefect/settings/models/runner.py +4 -4
  101. prefect/settings/models/server/api.py +3 -3
  102. prefect/settings/models/server/database.py +11 -4
  103. prefect/settings/models/server/deployments.py +6 -2
  104. prefect/settings/models/server/ephemeral.py +4 -2
  105. prefect/settings/models/server/events.py +3 -2
  106. prefect/settings/models/server/flow_run_graph.py +6 -2
  107. prefect/settings/models/server/root.py +3 -3
  108. prefect/settings/models/server/services.py +26 -11
  109. prefect/settings/models/server/tasks.py +6 -3
  110. prefect/settings/models/server/ui.py +3 -3
  111. prefect/settings/models/tasks.py +5 -5
  112. prefect/settings/models/testing.py +3 -3
  113. prefect/settings/models/worker.py +5 -3
  114. prefect/settings/profiles.py +15 -2
  115. prefect/states.py +61 -45
  116. prefect/task_engine.py +54 -75
  117. prefect/task_runners.py +56 -55
  118. prefect/task_worker.py +2 -2
  119. prefect/tasks.py +90 -36
  120. prefect/telemetry/bootstrap.py +10 -9
  121. prefect/telemetry/run_telemetry.py +13 -8
  122. prefect/telemetry/services.py +4 -0
  123. prefect/transactions.py +4 -15
  124. prefect/utilities/_git.py +34 -0
  125. prefect/utilities/asyncutils.py +1 -1
  126. prefect/utilities/engine.py +3 -19
  127. prefect/utilities/generics.py +18 -0
  128. prefect/utilities/templating.py +25 -1
  129. prefect/workers/base.py +6 -3
  130. prefect/workers/process.py +1 -1
  131. {prefect_client-3.1.10.dist-info → prefect_client-3.1.12.dist-info}/METADATA +2 -2
  132. {prefect_client-3.1.10.dist-info → prefect_client-3.1.12.dist-info}/RECORD +135 -109
  133. prefect/client/orchestration.py +0 -4523
  134. prefect/records/__init__.py +0 -1
  135. prefect/records/base.py +0 -235
  136. prefect/records/filesystem.py +0 -213
  137. prefect/records/memory.py +0 -184
  138. prefect/records/result_store.py +0 -70
  139. {prefect_client-3.1.10.dist-info → prefect_client-3.1.12.dist-info}/LICENSE +0 -0
  140. {prefect_client-3.1.10.dist-info → prefect_client-3.1.12.dist-info}/WHEEL +0 -0
  141. {prefect_client-3.1.10.dist-info → prefect_client-3.1.12.dist-info}/top_level.txt +0 -0
prefect/blocks/core.py CHANGED
@@ -1,3 +1,5 @@
1
+ from __future__ import annotations
2
+
1
3
  import hashlib
2
4
  import html
3
5
  import inspect
@@ -10,14 +12,10 @@ from textwrap import dedent
10
12
  from typing import (
11
13
  TYPE_CHECKING,
12
14
  Any,
13
- Callable,
14
15
  ClassVar,
15
- Dict,
16
+ Coroutine,
16
17
  FrozenSet,
17
- List,
18
18
  Optional,
19
- Tuple,
20
- Type,
21
19
  TypeVar,
22
20
  Union,
23
21
  get_origin,
@@ -34,11 +32,12 @@ from pydantic import (
34
32
  SecretBytes,
35
33
  SecretStr,
36
34
  SerializationInfo,
35
+ SerializerFunctionWrapHandler,
37
36
  ValidationError,
38
37
  model_serializer,
39
38
  )
40
39
  from pydantic.json_schema import GenerateJsonSchema
41
- from typing_extensions import Literal, ParamSpec, Self, get_args
40
+ from typing_extensions import Literal, ParamSpec, Self, TypeGuard, get_args
42
41
 
43
42
  import prefect.exceptions
44
43
  from prefect._internal.compatibility.async_dispatch import async_dispatch
@@ -49,6 +48,12 @@ from prefect.client.schemas import (
49
48
  BlockType,
50
49
  BlockTypeUpdate,
51
50
  )
51
+ from prefect.client.schemas.actions import (
52
+ BlockDocumentCreate,
53
+ BlockDocumentUpdate,
54
+ BlockSchemaCreate,
55
+ BlockTypeCreate,
56
+ )
52
57
  from prefect.client.utilities import inject_client
53
58
  from prefect.events import emit_event
54
59
  from prefect.logging.loggers import disable_logger
@@ -70,13 +75,15 @@ if TYPE_CHECKING:
70
75
  R = TypeVar("R")
71
76
  P = ParamSpec("P")
72
77
 
73
- ResourceTuple = Tuple[Dict[str, Any], List[Dict[str, Any]]]
78
+ ResourceTuple = tuple[dict[str, Any], list[dict[str, Any]]]
74
79
 
75
80
 
76
81
  def block_schema_to_key(schema: BlockSchema) -> str:
77
82
  """
78
83
  Defines the unique key used to lookup the Block class for a given schema.
79
84
  """
85
+ if schema.block_type is None:
86
+ raise ValueError("Block type is not set")
80
87
  return f"{schema.block_type.slug}"
81
88
 
82
89
 
@@ -87,14 +94,16 @@ class InvalidBlockRegistration(Exception):
87
94
  """
88
95
 
89
96
 
90
- def _collect_nested_reference_strings(obj: Dict) -> List[str]:
97
+ def _collect_nested_reference_strings(
98
+ obj: dict[str, Any] | list[Any],
99
+ ) -> list[dict[str, Any]]:
91
100
  """
92
101
  Collects all nested reference strings (e.g. #/definitions/Model) from a given object.
93
102
  """
94
- found_reference_strings = []
103
+ found_reference_strings: list[dict[str, Any]] = []
95
104
  if isinstance(obj, dict):
96
- if obj.get("$ref"):
97
- found_reference_strings.append(obj.get("$ref"))
105
+ if ref := obj.get("$ref"):
106
+ found_reference_strings.append(ref)
98
107
  for value in obj.values():
99
108
  found_reference_strings.extend(_collect_nested_reference_strings(value))
100
109
  if isinstance(obj, list):
@@ -103,28 +112,31 @@ def _collect_nested_reference_strings(obj: Dict) -> List[str]:
103
112
  return found_reference_strings
104
113
 
105
114
 
106
- def _get_non_block_reference_definitions(object_definition: Dict, definitions: Dict):
115
+ def _get_non_block_reference_definitions(
116
+ object_definition: dict[str, Any], definitions: dict[str, Any]
117
+ ) -> dict[str, Any]:
107
118
  """
108
119
  Given a definition of an object in a block schema OpenAPI spec and the dictionary
109
120
  of all reference definitions in that same block schema OpenAPI spec, return the
110
121
  definitions for objects that are referenced from the object or any children of
111
122
  the object that do not reference a block.
112
123
  """
113
- non_block_definitions = {}
124
+ non_block_definitions: dict[str, Any] = {}
114
125
  reference_strings = _collect_nested_reference_strings(object_definition)
115
126
  for reference_string in reference_strings:
116
- definition_key = reference_string.replace("#/definitions/", "")
117
- definition = definitions.get(definition_key)
118
- if definition and definition.get("block_type_slug") is None:
119
- non_block_definitions = {
120
- **non_block_definitions,
121
- definition_key: definition,
122
- **_get_non_block_reference_definitions(definition, definitions),
123
- }
127
+ if isinstance(reference_string, str):
128
+ definition_key = reference_string.replace("#/definitions/", "")
129
+ definition = definitions.get(definition_key)
130
+ if definition and definition.get("block_type_slug") is None:
131
+ non_block_definitions = {
132
+ **non_block_definitions,
133
+ definition_key: definition,
134
+ **_get_non_block_reference_definitions(definition, definitions),
135
+ }
124
136
  return non_block_definitions
125
137
 
126
138
 
127
- def _is_subclass(cls, parent_cls) -> bool:
139
+ def _is_subclass(cls: type, parent_cls: type) -> TypeGuard[type[BaseModel]]:
128
140
  """
129
141
  Checks if a given class is a subclass of another class. Unlike issubclass,
130
142
  this will not throw an exception if cls is an instance instead of a type.
@@ -135,7 +147,9 @@ def _is_subclass(cls, parent_cls) -> bool:
135
147
 
136
148
 
137
149
  def _collect_secret_fields(
138
- name: str, type_: Type[BaseModel], secrets: List[str]
150
+ name: str,
151
+ type_: type[BaseModel] | type[SecretStr] | type[SecretBytes] | type[SecretDict],
152
+ secrets: list[str],
139
153
  ) -> None:
140
154
  """
141
155
  Recursively collects all secret fields from a given type and adds them to the
@@ -148,7 +162,10 @@ def _collect_secret_fields(
148
162
  return
149
163
  elif _is_subclass(type_, BaseModel):
150
164
  for field_name, field in type_.model_fields.items():
151
- _collect_secret_fields(f"{name}.{field_name}", field.annotation, secrets)
165
+ if field.annotation is not None:
166
+ _collect_secret_fields(
167
+ f"{name}.{field_name}", field.annotation, secrets
168
+ )
152
169
  return
153
170
 
154
171
  if type_ in (SecretStr, SecretBytes) or (
@@ -212,7 +229,7 @@ class BlockNotSavedError(RuntimeError):
212
229
  pass
213
230
 
214
231
 
215
- def schema_extra(schema: Dict[str, Any], model: Type["Block"]):
232
+ def schema_extra(schema: dict[str, Any], model: type["Block"]) -> None:
216
233
  """
217
234
  Customizes Pydantic's schema generation feature to add blocks related information.
218
235
  """
@@ -231,31 +248,35 @@ def schema_extra(schema: Dict[str, Any], model: Type["Block"]):
231
248
  # for example: ["x", "y", "z.*", "child.a"]
232
249
  # means the top-level keys "x" and "y", all keys under "z", and the key "a" of a block
233
250
  # nested under the "child" key are all secret. There is no limit to nesting.
234
- secrets = schema["secret_fields"] = []
251
+ secrets: list[str] = []
235
252
  for name, field in model.model_fields.items():
236
- _collect_secret_fields(name, field.annotation, secrets)
253
+ if field.annotation is not None:
254
+ _collect_secret_fields(name, field.annotation, secrets)
255
+ schema["secret_fields"] = secrets
237
256
 
238
257
  # create block schema references
239
- refs = schema["block_schema_references"] = {}
258
+ refs: dict[str, Any] = {}
240
259
 
241
260
  def collect_block_schema_references(field_name: str, annotation: type) -> None:
242
261
  """Walk through the annotation and collect block schemas for any nested blocks."""
243
262
  if Block.is_block_class(annotation):
244
263
  if isinstance(refs.get(field_name), list):
245
- refs[field_name].append(annotation._to_block_schema_reference_dict())
264
+ refs[field_name].append(annotation._to_block_schema_reference_dict()) # pyright: ignore[reportPrivateUsage]
246
265
  elif isinstance(refs.get(field_name), dict):
247
266
  refs[field_name] = [
248
267
  refs[field_name],
249
- annotation._to_block_schema_reference_dict(),
268
+ annotation._to_block_schema_reference_dict(), # pyright: ignore[reportPrivateUsage]
250
269
  ]
251
270
  else:
252
- refs[field_name] = annotation._to_block_schema_reference_dict()
271
+ refs[field_name] = annotation._to_block_schema_reference_dict() # pyright: ignore[reportPrivateUsage]
253
272
  if get_origin(annotation) in (Union, list, tuple, dict):
254
273
  for type_ in get_args(annotation):
255
274
  collect_block_schema_references(field_name, type_)
256
275
 
257
276
  for name, field in model.model_fields.items():
258
- collect_block_schema_references(name, field.annotation)
277
+ if field.annotation is not None:
278
+ collect_block_schema_references(name, field.annotation)
279
+ schema["block_schema_references"] = refs
259
280
 
260
281
 
261
282
  @register_base_type
@@ -276,7 +297,7 @@ class Block(BaseModel, ABC):
276
297
  initialization.
277
298
  """
278
299
 
279
- model_config = ConfigDict(
300
+ model_config: ClassVar[ConfigDict] = ConfigDict(
280
301
  extra="allow",
281
302
  json_schema_extra=schema_extra,
282
303
  )
@@ -288,7 +309,7 @@ class Block(BaseModel, ABC):
288
309
  def __str__(self) -> str:
289
310
  return self.__repr__()
290
311
 
291
- def __repr_args__(self):
312
+ def __repr_args__(self) -> list[tuple[str | None, Any]]:
292
313
  repr_args = super().__repr_args__()
293
314
  data_keys = self.model_json_schema()["properties"].keys()
294
315
  return [
@@ -315,7 +336,7 @@ class Block(BaseModel, ABC):
315
336
  _code_example: ClassVar[Optional[str]] = None
316
337
  _block_type_id: ClassVar[Optional[UUID]] = None
317
338
  _block_schema_id: ClassVar[Optional[UUID]] = None
318
- _block_schema_capabilities: ClassVar[Optional[List[str]]] = None
339
+ _block_schema_capabilities: ClassVar[Optional[list[str]]] = None
319
340
  _block_schema_version: ClassVar[Optional[str]] = None
320
341
 
321
342
  # -- private instance variables
@@ -327,18 +348,20 @@ class Block(BaseModel, ABC):
327
348
 
328
349
  # Exclude `save` as it uses the `sync_compatible` decorator and needs to be
329
350
  # decorated directly.
330
- _events_excluded_methods: ClassVar[List[str]] = PrivateAttr(
351
+ _events_excluded_methods: ClassVar[list[str]] = PrivateAttr(
331
352
  default=["block_initialization", "save", "dict"]
332
353
  )
333
354
 
334
355
  @classmethod
335
- def __dispatch_key__(cls):
356
+ def __dispatch_key__(cls) -> str | None:
336
357
  if cls.__name__ == "Block":
337
358
  return None # The base class is abstract
338
359
  return block_schema_to_key(cls._to_block_schema())
339
360
 
340
361
  @model_serializer(mode="wrap")
341
- def ser_model(self, handler: Callable, info: SerializationInfo) -> Any:
362
+ def ser_model(
363
+ self, handler: SerializerFunctionWrapHandler, info: SerializationInfo
364
+ ) -> Any:
342
365
  jsonable_self = handler(self)
343
366
  if (ctx := info.context) and ctx.get("include_secrets") is True:
344
367
  jsonable_self.update(
@@ -363,11 +386,11 @@ class Block(BaseModel, ABC):
363
386
  return jsonable_self
364
387
 
365
388
  @classmethod
366
- def get_block_type_name(cls):
389
+ def get_block_type_name(cls) -> str:
367
390
  return cls._block_type_name or cls.__name__
368
391
 
369
392
  @classmethod
370
- def get_block_type_slug(cls):
393
+ def get_block_type_slug(cls) -> str:
371
394
  return slugify(cls._block_type_slug or cls.get_block_type_name())
372
395
 
373
396
  @classmethod
@@ -413,7 +436,7 @@ class Block(BaseModel, ABC):
413
436
 
414
437
  @classmethod
415
438
  def _calculate_schema_checksum(
416
- cls, block_schema_fields: Optional[Dict[str, Any]] = None
439
+ cls, block_schema_fields: dict[str, Any] | None = None
417
440
  ):
418
441
  """
419
442
  Generates a unique hash for the underlying schema of block.
@@ -513,11 +536,24 @@ class Block(BaseModel, ABC):
513
536
  "$ref": {"block_document_id": field_value._block_document_id}
514
537
  }
515
538
 
539
+ block_schema_id = block_schema_id or self._block_schema_id
540
+ block_type_id = block_type_id or self._block_type_id
541
+
542
+ if block_schema_id is None:
543
+ raise ValueError(
544
+ "No block schema ID provided, either as an argument or on the block."
545
+ )
546
+ if block_type_id is None:
547
+ raise ValueError(
548
+ "No block type ID provided, either as an argument or on the block."
549
+ )
550
+
516
551
  return BlockDocument(
517
552
  id=self._block_document_id or uuid4(),
518
553
  name=(name or self._block_document_name) if not is_anonymous else None,
519
- block_schema_id=block_schema_id or self._block_schema_id,
520
- block_type_id=block_type_id or self._block_type_id,
554
+ block_schema_id=block_schema_id,
555
+ block_type_id=block_type_id,
556
+ block_type_name=self._block_type_name,
521
557
  data=block_document_data,
522
558
  block_schema=self._to_block_schema(
523
559
  block_type_id=block_type_id or self._block_type_id,
@@ -551,13 +587,15 @@ class Block(BaseModel, ABC):
551
587
  )
552
588
 
553
589
  @classmethod
554
- def _parse_docstring(cls) -> List[DocstringSection]:
590
+ def _parse_docstring(cls) -> list[DocstringSection]:
555
591
  """
556
592
  Parses the docstring into list of DocstringSection objects.
557
593
  Helper method used primarily to suppress irrelevant logs, e.g.
558
594
  `<module>:11: No type or annotation for parameter 'write_json'`
559
595
  because griffe is unable to parse the types from pydantic.BaseModel.
560
596
  """
597
+ if cls.__doc__ is None:
598
+ return []
561
599
  with disable_logger("griffe"):
562
600
  docstring = Docstring(cls.__doc__)
563
601
  parsed = parse(docstring, Parser.google)
@@ -710,7 +748,7 @@ class Block(BaseModel, ABC):
710
748
  def _event_kind(self) -> str:
711
749
  return f"prefect.block.{self.get_block_type_slug()}"
712
750
 
713
- def _event_method_called_resources(self) -> Optional[ResourceTuple]:
751
+ def _event_method_called_resources(self) -> ResourceTuple | None:
714
752
  if not (self._block_document_id and self._block_document_name):
715
753
  return None
716
754
 
@@ -732,14 +770,14 @@ class Block(BaseModel, ABC):
732
770
  )
733
771
 
734
772
  @classmethod
735
- def get_block_class_from_schema(cls: Type[Self], schema: BlockSchema) -> Type[Self]:
773
+ def get_block_class_from_schema(cls: type[Self], schema: BlockSchema) -> type[Self]:
736
774
  """
737
775
  Retrieve the block class implementation given a schema.
738
776
  """
739
777
  return cls.get_block_class_from_key(block_schema_to_key(schema))
740
778
 
741
779
  @classmethod
742
- def get_block_class_from_key(cls: Type[Self], key: str) -> Type[Self]:
780
+ def get_block_class_from_key(cls: type[Self], key: str) -> type[Self]:
743
781
  """
744
782
  Retrieve the block class implementation given a key.
745
783
  """
@@ -751,7 +789,7 @@ class Block(BaseModel, ABC):
751
789
  return lookup_type(cls, key)
752
790
 
753
791
  def _define_metadata_on_nested_blocks(
754
- self, block_document_references: Dict[str, Dict[str, Any]]
792
+ self, block_document_references: dict[str, dict[str, Any]]
755
793
  ):
756
794
  """
757
795
  Recursively populates metadata fields on nested blocks based on the
@@ -827,13 +865,14 @@ class Block(BaseModel, ABC):
827
865
  return block_document, block_document_name
828
866
 
829
867
  @classmethod
830
- @sync_compatible
831
868
  @inject_client
832
869
  async def _get_block_document_by_id(
833
870
  cls,
834
871
  block_document_id: Union[str, uuid.UUID],
835
- client: Optional["PrefectClient"] = None,
872
+ client: "PrefectClient | None" = None,
836
873
  ):
874
+ if TYPE_CHECKING:
875
+ assert isinstance(client, PrefectClient)
837
876
  if isinstance(block_document_id, str):
838
877
  try:
839
878
  block_document_id = UUID(block_document_id)
@@ -1038,7 +1077,6 @@ class Block(BaseModel, ABC):
1038
1077
  block_document, _ = run_coro_as_sync(
1039
1078
  cls._aget_block_document(name, client=client)
1040
1079
  )
1041
-
1042
1080
  return cls._load_from_block_document(block_document, validate=validate)
1043
1081
 
1044
1082
  @classmethod
@@ -1046,10 +1084,10 @@ class Block(BaseModel, ABC):
1046
1084
  @inject_client
1047
1085
  async def load_from_ref(
1048
1086
  cls,
1049
- ref: Union[str, UUID, Dict[str, Any]],
1087
+ ref: Union[str, UUID, dict[str, Any]],
1050
1088
  validate: bool = True,
1051
- client: Optional["PrefectClient"] = None,
1052
- ) -> "Self":
1089
+ client: "PrefectClient | None" = None,
1090
+ ) -> Self:
1053
1091
  """
1054
1092
  Retrieves data from the block document by given reference for the block type
1055
1093
  that corresponds with the current class and returns an instantiated version of
@@ -1089,16 +1127,18 @@ class Block(BaseModel, ABC):
1089
1127
  block document with the specified name.
1090
1128
 
1091
1129
  """
1130
+ if TYPE_CHECKING:
1131
+ assert isinstance(client, PrefectClient)
1092
1132
  block_document = None
1093
1133
  if isinstance(ref, (str, UUID)):
1094
1134
  block_document, _ = await cls._get_block_document_by_id(ref, client=client)
1095
- elif isinstance(ref, dict):
1135
+ else:
1096
1136
  if block_document_id := ref.get("block_document_id"):
1097
1137
  block_document, _ = await cls._get_block_document_by_id(
1098
1138
  block_document_id, client=client
1099
1139
  )
1100
1140
  elif block_document_slug := ref.get("block_document_slug"):
1101
- block_document, _ = await cls._get_block_document(
1141
+ block_document, _ = await cls._aget_block_document(
1102
1142
  block_document_slug, client=client
1103
1143
  )
1104
1144
 
@@ -1110,7 +1150,7 @@ class Block(BaseModel, ABC):
1110
1150
  @classmethod
1111
1151
  def _load_from_block_document(
1112
1152
  cls, block_document: BlockDocument, validate: bool = True
1113
- ) -> "Self":
1153
+ ) -> Self:
1114
1154
  """
1115
1155
  Loads a block from a given block document.
1116
1156
 
@@ -1144,7 +1184,9 @@ class Block(BaseModel, ABC):
1144
1184
  except ValidationError as e:
1145
1185
  if not validate:
1146
1186
  missing_fields = tuple(err["loc"][0] for err in e.errors())
1147
- missing_block_data = {field: None for field in missing_fields}
1187
+ missing_block_data: dict[str, None] = {
1188
+ field: None for field in missing_fields if isinstance(field, str)
1189
+ }
1148
1190
  warnings.warn(
1149
1191
  f"Could not fully load {block_document.name!r} of block type"
1150
1192
  f" {cls.get_block_type_slug()!r} - this is likely because one or more"
@@ -1163,7 +1205,7 @@ class Block(BaseModel, ABC):
1163
1205
  ) from e
1164
1206
 
1165
1207
  @staticmethod
1166
- def is_block_class(block) -> bool:
1208
+ def is_block_class(block: Any) -> TypeGuard[type["Block"]]:
1167
1209
  return _is_subclass(block, Block)
1168
1210
 
1169
1211
  @staticmethod
@@ -1191,6 +1233,8 @@ class Block(BaseModel, ABC):
1191
1233
  Prefect API. A new client will be created and used if one is not
1192
1234
  provided.
1193
1235
  """
1236
+ if TYPE_CHECKING:
1237
+ assert isinstance(client, PrefectClient)
1194
1238
  if cls.__name__ == "Block":
1195
1239
  raise InvalidBlockRegistration(
1196
1240
  "`register_type_and_schema` should be called on a Block "
@@ -1205,13 +1249,17 @@ class Block(BaseModel, ABC):
1205
1249
  async def register_blocks_in_annotation(annotation: type) -> None:
1206
1250
  """Walk through the annotation and register any nested blocks."""
1207
1251
  if Block.is_block_class(annotation):
1208
- await annotation.register_type_and_schema(client=client)
1252
+ coro = annotation.register_type_and_schema(client=client)
1253
+ if TYPE_CHECKING:
1254
+ assert isinstance(coro, Coroutine)
1255
+ await coro
1209
1256
  elif get_origin(annotation) in (Union, tuple, list, dict):
1210
1257
  for inner_annotation in get_args(annotation):
1211
1258
  await register_blocks_in_annotation(inner_annotation)
1212
1259
 
1213
1260
  for field in cls.model_fields.values():
1214
- await register_blocks_in_annotation(field.annotation)
1261
+ if field.annotation is not None:
1262
+ await register_blocks_in_annotation(field.annotation)
1215
1263
 
1216
1264
  try:
1217
1265
  block_type = await client.read_block_type_by_slug(
@@ -1225,10 +1273,32 @@ class Block(BaseModel, ABC):
1225
1273
  local_block_type=local_block_type, server_block_type=block_type
1226
1274
  ):
1227
1275
  await client.update_block_type(
1228
- block_type_id=block_type.id, block_type=local_block_type
1276
+ block_type_id=block_type.id,
1277
+ block_type=BlockTypeUpdate(
1278
+ **local_block_type.model_dump(
1279
+ include={
1280
+ "logo_url",
1281
+ "documentation_url",
1282
+ "description",
1283
+ "code_example",
1284
+ }
1285
+ )
1286
+ ),
1229
1287
  )
1230
1288
  except prefect.exceptions.ObjectNotFound:
1231
- block_type = await client.create_block_type(block_type=cls._to_block_type())
1289
+ block_type_create = BlockTypeCreate(
1290
+ **cls._to_block_type().model_dump(
1291
+ include={
1292
+ "name",
1293
+ "slug",
1294
+ "logo_url",
1295
+ "documentation_url",
1296
+ "description",
1297
+ "code_example",
1298
+ }
1299
+ )
1300
+ )
1301
+ block_type = await client.create_block_type(block_type=block_type_create)
1232
1302
  cls._block_type_id = block_type.id
1233
1303
 
1234
1304
  try:
@@ -1237,8 +1307,13 @@ class Block(BaseModel, ABC):
1237
1307
  version=cls.get_block_schema_version(),
1238
1308
  )
1239
1309
  except prefect.exceptions.ObjectNotFound:
1310
+ block_schema_create = BlockSchemaCreate(
1311
+ **cls._to_block_schema(block_type_id=block_type.id).model_dump(
1312
+ include={"fields", "block_type_id", "capabilities", "version"}
1313
+ )
1314
+ )
1240
1315
  block_schema = await client.create_block_schema(
1241
- block_schema=cls._to_block_schema(block_type_id=block_type.id)
1316
+ block_schema=block_schema_create
1242
1317
  )
1243
1318
 
1244
1319
  cls._block_schema_id = block_schema.id
@@ -1250,7 +1325,7 @@ class Block(BaseModel, ABC):
1250
1325
  is_anonymous: bool = False,
1251
1326
  overwrite: bool = False,
1252
1327
  client: Optional["PrefectClient"] = None,
1253
- ):
1328
+ ) -> UUID:
1254
1329
  """
1255
1330
  Saves the values of a block as a block document with an option to save as an
1256
1331
  anonymous block document.
@@ -1268,6 +1343,8 @@ class Block(BaseModel, ABC):
1268
1343
  ValueError: If a name is not given and `is_anonymous` is `False` or a name is given and
1269
1344
  `is_anonymous` is `True`.
1270
1345
  """
1346
+ if TYPE_CHECKING:
1347
+ assert isinstance(client, PrefectClient)
1271
1348
  if name is None and not is_anonymous:
1272
1349
  if self._block_document_name is None:
1273
1350
  raise ValueError(
@@ -1281,25 +1358,47 @@ class Block(BaseModel, ABC):
1281
1358
  self._is_anonymous = is_anonymous
1282
1359
 
1283
1360
  # Ensure block type and schema are registered before saving block document.
1284
- await self.register_type_and_schema(client=client)
1361
+ coro = self.register_type_and_schema(client=client)
1362
+ if TYPE_CHECKING:
1363
+ assert isinstance(coro, Coroutine)
1364
+ await coro
1285
1365
 
1366
+ block_document = None
1286
1367
  try:
1368
+ block_document_create = BlockDocumentCreate(
1369
+ **self._to_block_document(name=name, include_secrets=True).model_dump(
1370
+ include={
1371
+ "name",
1372
+ "block_schema_id",
1373
+ "block_type_id",
1374
+ "data",
1375
+ "is_anonymous",
1376
+ }
1377
+ )
1378
+ )
1287
1379
  block_document = await client.create_block_document(
1288
- block_document=self._to_block_document(name=name)
1380
+ block_document=block_document_create
1289
1381
  )
1290
1382
  except prefect.exceptions.ObjectAlreadyExists as err:
1291
1383
  if overwrite:
1292
1384
  block_document_id = self._block_document_id
1293
- if block_document_id is None:
1385
+ if block_document_id is None and name is not None:
1294
1386
  existing_block_document = await client.read_block_document_by_name(
1295
1387
  name=name, block_type_slug=self.get_block_type_slug()
1296
1388
  )
1297
1389
  block_document_id = existing_block_document.id
1390
+ if TYPE_CHECKING:
1391
+ # We know that the block document id is not None here because we
1392
+ # only get here if the block document already exists
1393
+ assert isinstance(block_document_id, UUID)
1394
+ block_document_update = BlockDocumentUpdate(
1395
+ **self._to_block_document(
1396
+ name=name, include_secrets=True
1397
+ ).model_dump(include={"block_schema_id", "data"})
1398
+ )
1298
1399
  await client.update_block_document(
1299
1400
  block_document_id=block_document_id,
1300
- block_document=self._to_block_document(
1301
- name=name, include_secrets=True
1302
- ),
1401
+ block_document=block_document_update,
1303
1402
  )
1304
1403
  block_document = await client.read_block_document(
1305
1404
  block_document_id=block_document_id
@@ -1351,7 +1450,7 @@ class Block(BaseModel, ABC):
1351
1450
 
1352
1451
  await client.delete_block_document(block_document.id)
1353
1452
 
1354
- def __new__(cls: Type[Self], **kwargs) -> Self:
1453
+ def __new__(cls: type[Self], **kwargs: Any) -> Self:
1355
1454
  """
1356
1455
  Create an instance of the Block subclass type if a `block_type_slug` is
1357
1456
  present in the data payload.
@@ -1390,9 +1489,9 @@ class Block(BaseModel, ABC):
1390
1489
  cls,
1391
1490
  by_alias: bool = True,
1392
1491
  ref_template: str = "#/definitions/{model}",
1393
- schema_generator: Type[GenerateJsonSchema] = GenerateJsonSchema,
1492
+ schema_generator: type[GenerateJsonSchema] = GenerateJsonSchema,
1394
1493
  mode: Literal["validation", "serialization"] = "validation",
1395
- ) -> Dict[str, Any]:
1494
+ ) -> dict[str, Any]:
1396
1495
  """TODO: stop overriding this method - use GenerateSchema in ConfigDict instead?"""
1397
1496
  schema = super().model_json_schema(
1398
1497
  by_alias, ref_template, schema_generator, mode
@@ -1415,11 +1514,11 @@ class Block(BaseModel, ABC):
1415
1514
  @classmethod
1416
1515
  def model_validate(
1417
1516
  cls: type[Self],
1418
- obj: Any,
1517
+ obj: dict[str, Any] | Any,
1419
1518
  *,
1420
- strict: Optional[bool] = None,
1421
- from_attributes: Optional[bool] = None,
1422
- context: Optional[Dict[str, Any]] = None,
1519
+ strict: bool | None = None,
1520
+ from_attributes: bool | None = None,
1521
+ context: dict[str, Any] | None = None,
1423
1522
  ) -> Self:
1424
1523
  if isinstance(obj, dict):
1425
1524
  extra_serializer_fields = {
@@ -1437,18 +1536,18 @@ class Block(BaseModel, ABC):
1437
1536
  def model_dump(
1438
1537
  self,
1439
1538
  *,
1440
- mode: Union[Literal["json", "python"], str] = "python",
1441
- include: "IncEx" = None,
1442
- exclude: "IncEx" = None,
1443
- context: Optional[Dict[str, Any]] = None,
1539
+ mode: Literal["json", "python"] | str = "python",
1540
+ include: "IncEx | None" = None,
1541
+ exclude: "IncEx | None" = None,
1542
+ context: dict[str, Any] | None = None,
1444
1543
  by_alias: bool = False,
1445
1544
  exclude_unset: bool = False,
1446
1545
  exclude_defaults: bool = False,
1447
1546
  exclude_none: bool = False,
1448
1547
  round_trip: bool = False,
1449
- warnings: Union[bool, Literal["none", "warn", "error"]] = True,
1548
+ warnings: bool | Literal["none", "warn", "error"] = True,
1450
1549
  serialize_as_any: bool = False,
1451
- ) -> Dict[str, Any]:
1550
+ ) -> dict[str, Any]:
1452
1551
  d = super().model_dump(
1453
1552
  mode=mode,
1454
1553
  include=include,