Coverage for /usr/local/lib/python3.12/site-packages/prefect/_internal/schemas/bases.py: 51%

47 statements  

« prev     ^ index     » next       coverage.py v7.10.6, created at 2025-12-05 10:48 +0000

1""" 

2Utilities for creating and working with Prefect REST API schemas. 

3""" 

4 

5import datetime 1a

6from typing import Any, ClassVar, Optional, TypeVar 1a

7from uuid import UUID, uuid4 1a

8 

9from pydantic import BaseModel, ConfigDict, Field 1a

10from rich.repr import RichReprResult 1a

11from typing_extensions import Self 1a

12 

13from prefect._internal.uuid7 import uuid7 1a

14from prefect.types._datetime import ( 1a

15 DateTime, 

16 human_friendly_diff, 

17) 

18from prefect.utilities.generics import validate_list 1a

19 

20T = TypeVar("T") 1a

21 

22 

23class PrefectBaseModel(BaseModel): 1a

24 """A base pydantic.BaseModel for all Prefect schemas and pydantic models. 

25 

26 As the basis for most Prefect schemas, this base model usually ignores extra 

27 fields that are passed to it at instantiation. Because adding new fields to 

28 API payloads is not considered a breaking change, this ensures that any 

29 Prefect client loading data from a server running a possibly-newer version 

30 of Prefect will be able to process those new fields gracefully. 

31 """ 

32 

33 _reset_fields: ClassVar[set[str]] = set() 1a

34 

35 model_config: ClassVar[ConfigDict] = ConfigDict( 1a

36 ser_json_timedelta="float", 

37 defer_build=True, 

38 extra="ignore", 

39 ) 

40 

41 def __eq__(self, other: Any) -> bool: 1a

42 """Equality operator that ignores the resettable fields of the PrefectBaseModel. 

43 

44 NOTE: this equality operator will only be applied if the PrefectBaseModel is 

45 the left-hand operand. This is a limitation of Python. 

46 """ 

47 copy_dict = self.model_dump(exclude=self._reset_fields) 

48 if isinstance(other, PrefectBaseModel): 

49 return copy_dict == other.model_dump(exclude=other._reset_fields) 

50 if isinstance(other, BaseModel): 

51 return copy_dict == other.model_dump() 

52 else: 

53 return copy_dict == other 

54 

55 @classmethod 1a

56 def model_validate_list( 1a

57 cls, 

58 obj: Any, 

59 *, 

60 strict: Optional[bool] = None, 

61 from_attributes: Optional[bool] = None, 

62 context: Optional[Any] = None, 

63 ) -> list[Self]: 

64 return validate_list(cls, obj) 

65 

66 def __rich_repr__(self) -> RichReprResult: 1a

67 # Display all of the fields in the model if they differ from the default value 

68 for name, field in type(self).model_fields.items(): 

69 value = getattr(self, name) 

70 

71 # Simplify the display of some common fields 

72 if field.annotation == UUID and value: 

73 value = str(value) 

74 elif ( 

75 isinstance(field.annotation, datetime.datetime) 

76 and name == "timestamp" 

77 and value 

78 ): 

79 value = value.isoformat() 

80 elif isinstance(field.annotation, datetime.datetime) and value: 

81 value = human_friendly_diff(value) 

82 

83 yield name, value, field.get_default() 

84 

85 def reset_fields(self: Self) -> Self: 1a

86 """ 

87 Reset the fields of the model that are in the `_reset_fields` set. 

88 

89 Returns: 

90 PrefectBaseModel: A new instance of the model with the reset fields. 

91 """ 

92 return self.model_copy( 

93 update={ 

94 field: type(self) 

95 .model_fields[field] 

96 .get_default(call_default_factory=True) 

97 for field in self._reset_fields 

98 } 

99 ) 

100 

101 

102class IDBaseModel(PrefectBaseModel): 1a

103 """ 

104 A PrefectBaseModel with a randomly-generated UUID ID value. 

105 

106 The ID is reset on copy() and not included in equality comparisons. 

107 """ 

108 

109 _reset_fields: ClassVar[set[str]] = {"id"} 1a

110 id: UUID = Field(default_factory=uuid4) 1a

111 

112 

113class TimeSeriesBaseModel(IDBaseModel): 1a

114 """ 

115 A PrefectBaseModel with a time-oriented UUIDv7 ID value. Used for models that 

116 operate like timeseries, such as runs, states, and logs. 

117 """ 

118 

119 id: UUID = Field(default_factory=uuid7) 1a

120 

121 

122class ObjectBaseModel(IDBaseModel): 1a

123 """ 

124 A PrefectBaseModel with an auto-generated UUID ID value and created / 

125 updated timestamps, intended for compatibility with our standard ORM models. 

126 

127 The ID, created, and updated fields are reset on copy() and not included in 

128 equality comparisons. 

129 """ 

130 

131 _reset_fields: ClassVar[set[str]] = {"id", "created", "updated"} 1a

132 model_config: ClassVar[ConfigDict] = ConfigDict(from_attributes=True) 1a

133 

134 created: Optional[DateTime] = Field(default=None, repr=False) 1a

135 updated: Optional[DateTime] = Field(default=None, repr=False) 1a

136 

137 

138class ActionBaseModel(PrefectBaseModel): 1a

139 model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") 1a