Coverage for polar/kit/schemas.py: 89%

1import dataclasses 

2import json 

3from collections.abc import Sequence 

4from datetime import datetime 

5from typing import Annotated, Any, Literal, cast, get_args, overload 

6 

7from pydantic import ( 

8 UUID4, 

9 AfterValidator, 

10 BaseModel, 

11 ConfigDict, 

12 Field, 

13 GetCoreSchemaHandler, 

14 GetJsonSchemaHandler, 

15 HttpUrl, 

16 PlainSerializer, 

17) 

18from pydantic.json_schema import JsonSchemaValue 

19from pydantic_core import CoreSchema, core_schema 

20from slugify import slugify 

21 

22 

23class Schema(BaseModel): 

24 model_config = ConfigDict(from_attributes=True) 

25 

26 

27class IDSchema(Schema): 

28 id: UUID4 = Field(..., description="The ID of the object.") 

29 

30 model_config = ConfigDict( 

31 # IMPORTANT: this ensures FastAPI doesn't generate `-Input` for output schemas 

32 json_schema_mode_override="serialization", 

33 ) 

34 

35 

36class TimestampedSchema(Schema): 

37 created_at: datetime = Field(description="Creation timestamp of the object.") 

38 modified_at: datetime | None = Field( 

39 description="Last modification timestamp of the object." 

40 ) 

41 

42 

43def empty_str_to_none(value: str | None) -> str | None: 

44 if isinstance(value, str): 

45 stripped_value = value.strip() 

46 if stripped_value == "": 

47 return None 

48 return stripped_value 

49 return value 

50 

51 

52EmptyStrToNoneValidator = AfterValidator(empty_str_to_none) 

53EmptyStrToNone = Annotated[str | None, EmptyStrToNoneValidator] 

54 

55 

56def _validate_slug(value: str) -> str: 

57 slugified = slugify(value) 

58 if slugified != value: 

59 raise ValueError( 

60 "The slug can only contain ASCII letters, numbers and hyphens." 

61 ) 

62 return value 

63 

64 

65SlugValidator = AfterValidator(_validate_slug) 

66 

67UUID4ToStr = Annotated[UUID4, PlainSerializer(lambda v: str(v), return_type=str)] 

68HttpUrlToStr = Annotated[HttpUrl, PlainSerializer(lambda v: str(v), return_type=str)] 

69 

70 

71@dataclasses.dataclass(slots=True) 

72class ClassName: 

73 """ 

74 Used as an annotation metadata, it allows us to customize the name generated 

75 by Pydantic for a type; in particular, a long union. 

76 

77 It does **nothing** on its own, but it can be used by other classes. 

78 

79 Currently, it's used by `ListResource` to generate a shorter name for the 

80 OpenAPI schema, when we list a resource having a long union type. 

81 """ 

82 

83 name: str 

84 

85 def __hash__(self) -> int: 

86 return hash(type(self.name)) 

87 

88 

89@dataclasses.dataclass(slots=True) 

90class MergeJSONSchema: 

91 json_schema: JsonSchemaValue 

92 mode: Literal["validation", "serialization"] | None = None 

93 

94 def __get_pydantic_json_schema__( 

95 self, core_schema: core_schema.CoreSchema, handler: GetJsonSchemaHandler 

96 ) -> JsonSchemaValue: 

97 mode = self.mode or handler.mode 

98 json_schema = handler(core_schema) 

99 if mode != handler.mode: 99 ↛ 100line 99 didn't jump to line 100 because the condition on line 99 was never true

100 return json_schema 

101 return {**json_schema, **self.json_schema} 

102 

103 def __hash__(self) -> int: 

104 return hash(type(self.mode)) 

105 

106 

107@dataclasses.dataclass(slots=True) 

108class SetSchemaReference: 

109 ref_name: str 

110 

111 def __get_pydantic_core_schema__( 

112 self, source_type: Any, handler: GetCoreSchemaHandler 

113 ) -> CoreSchema: 

114 schema = handler(source_type) 

115 schema["ref"] = self.ref_name # type: ignore 

116 return schema 

117 

118 def __hash__(self) -> int: 

119 return hash(type(self.ref_name)) 

120 

121 

122@dataclasses.dataclass(slots=True) 

123class SelectorWidget: 

124 resource_root: str 

125 resource_name: str 

126 display_property: str 

127 

128 def __get_pydantic_json_schema__( 

129 self, core_schema: core_schema.CoreSchema, handler: GetJsonSchemaHandler 

130 ) -> JsonSchemaValue: 

131 json_schema = handler(core_schema) 

132 return {**json_schema, **self._get_extra_attributes()} 

133 

134 def _get_extra_attributes(self) -> dict[str, Any]: 

135 return { 

136 "x-polar-selector-widget": { 

137 "resourceRoot": self.resource_root, 

138 "resourceName": self.resource_name, 

139 "displayProperty": self.display_property, 

140 } 

141 } 

142 

143 def __hash__(self) -> int: 

144 return hash(json.dumps(self._get_extra_attributes())) 

145 

146 

147class MultipleQueryFilter[Q](Sequence[Q]): 

148 """ 

149 Custom type to handle query filters that can be either 

150 a single value or a list of values. 

151 

152 By customizing the schema generation, we can make it accept 

153 either a scalar or a list of values for the query parameter. 

154 

155 At runtime, we make sure that the value is always a list. 

156 """ 

157 

158 def __init__(self, v: Sequence[Q]): 

159 self.v = v 

160 

161 @overload 

162 def __getitem__(self, s: int) -> Q: ... 162 ↛ anywhereline 162 didn't jump anywhere: it always raised an exception.

163 

164 @overload 

165 def __getitem__(self, s: slice) -> Sequence[Q]: ... 165 ↛ anywhereline 165 didn't jump anywhere: it always raised an exception.

166 

167 def __getitem__(self, s: int | slice) -> Q | Sequence[Q]: 

168 return self.v[s] 

169 

170 def __len__(self) -> int: 

171 return len(self.v) 

172 

173 @classmethod 

174 def __get_pydantic_core_schema__( 

175 cls, source: Any, handler: GetCoreSchemaHandler 

176 ) -> core_schema.CoreSchema: 

177 args = get_args(source) 

178 if len(args) == 0: 178 ↛ 179line 178 didn't jump to line 179 because the condition on line 178 was never true

179 raise TypeError("QueryFilter requires at least one type argument") 

180 

181 generic_type = args[0] 

182 sequence_schema = handler.generate_schema(Sequence[generic_type]) # type: ignore 

183 scalar_schema = handler.generate_schema(generic_type) 

184 union_schema = core_schema.union_schema([scalar_schema, sequence_schema]) 

185 

186 return core_schema.no_info_after_validator_function( 

187 cls._scalar_to_sequence, handler(union_schema) 

188 ) 

189 

190 @classmethod 

191 def _scalar_to_sequence(cls, v: Q | Sequence[Q]) -> Sequence[Q]: 

192 if isinstance(v, Sequence) and not isinstance(v, str): 192 ↛ 194line 192 didn't jump to line 194 because the condition on line 192 was always true

193 return v 

194 return [cast(Q, v)] # type: ignore[redundant-cast] 

195 

196 

197ORGANIZATION_ID_EXAMPLE = "1dbfc517-0bbf-4301-9ba8-555ca42b9737" 

198PRODUCT_ID_EXAMPLE = "d8dd2de1-21b7-4a41-8bc3-ce909c0cfe23" 

199PRICE_ID_EXAMPLE = "196ca717-4d84-4d28-a1b8-777255797dbc" 

200BENEFIT_ID_EXAMPLE = "397a17aa-15cf-4cb4-9333-18040203cf98" 

201CUSTOMER_ID_EXAMPLE = "992fae2a-2a17-4b7a-8d9e-e287cf90131b" 

202SUBSCRIPTION_ID_EXAMPLE = "e5149aae-e521-42b9-b24c-abb3d71eea2e" 

203BENEFIT_GRANT_ID_EXAMPLE = "d322132c-a9d0-4e0d-b8d3-d81ad021a3a9" 

204METER_ID_EXAMPLE = "d498a884-e2cd-4d3e-8002-f536468a8b22" 

205CHECKOUT_ID_EXAMPLE = "e4b478fa-cd25-4253-9f1f-8a41e6370ede"