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

1import dataclasses 2 ctx1ba

2import json 2 ctx1ba

3from collections.abc import Sequence 2 ctx1ba

4from datetime import datetime 2 ctx1ba

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

6 

7from pydantic import ( 2 ctx1ba

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 2 ctx1ba

19from pydantic_core import CoreSchema, core_schema 2 ctx1ba

20from slugify import slugify 2 ctx1ba

21 

22 

23class Schema(BaseModel): 2 ctx1ba

24 model_config = ConfigDict(from_attributes=True) 2 ctx1ba

25 

26 

27class IDSchema(Schema): 2 ctx1ba

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

29 

30 model_config = ConfigDict( 2 ctx1ba

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

32 json_schema_mode_override="serialization", 

33 ) 

34 

35 

36class TimestampedSchema(Schema): 2 ctx1ba

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

38 modified_at: datetime | None = Field( 2 ctx1ba

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

40 ) 

41 

42 

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

44 if isinstance(value, str): 44 ↛ 49line 44 didn't jump to line 49 because the condition on line 44 was always true3 ctx1bac

45 stripped_value = value.strip() 3 ctx1bac

46 if stripped_value == "": 46 ↛ 47line 46 didn't jump to line 47 because the condition on line 46 was never true3 ctx1bac

47 return None 

48 return stripped_value 3 ctx1bac

49 return value 

50 

51 

52EmptyStrToNoneValidator = AfterValidator(empty_str_to_none) 2 ctx1ba

53EmptyStrToNone = Annotated[str | None, EmptyStrToNoneValidator] 2 ctx1ba

54 

55 

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

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) 2 ctx1ba

66 

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

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

69 

70 

71@dataclasses.dataclass(slots=True) 2 ctx1ba

72class ClassName: 2 ctx1ba

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 2 ctx1ba

84 

85 def __hash__(self) -> int: 2 ctx1ba

86 return hash(type(self.name)) 1 ctx1a

87 

88 

89@dataclasses.dataclass(slots=True) 2 ctx1ba

90class MergeJSONSchema: 2 ctx1ba

91 json_schema: JsonSchemaValue 2 ctx1ba

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

93 

94 def __get_pydantic_json_schema__( 2 ctx1ba

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

96 ) -> JsonSchemaValue: 

97 mode = self.mode or handler.mode 1 ctx1c

98 json_schema = handler(core_schema) 1 ctx1c

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

100 return json_schema 

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

102 

103 def __hash__(self) -> int: 2 ctx1ba

104 return hash(type(self.mode)) 1 ctx1a

105 

106 

107@dataclasses.dataclass(slots=True) 2 ctx1ba

108class SetSchemaReference: 2 ctx1ba

109 ref_name: str 2 ctx1ba

110 

111 def __get_pydantic_core_schema__( 2 ctx1ba

112 self, source_type: Any, handler: GetCoreSchemaHandler 

113 ) -> CoreSchema: 

114 schema = handler(source_type) 3 ctx1acd

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

116 return schema 3 ctx1acd

117 

118 def __hash__(self) -> int: 2 ctx1ba

119 return hash(type(self.ref_name)) 1 ctx1a

120 

121 

122@dataclasses.dataclass(slots=True) 2 ctx1ba

123class SelectorWidget: 2 ctx1ba

124 resource_root: str 2 ctx1ba

125 resource_name: str 2 ctx1ba

126 display_property: str 2 ctx1ba

127 

128 def __get_pydantic_json_schema__( 2 ctx1ba

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

130 ) -> JsonSchemaValue: 

131 json_schema = handler(core_schema) 1 ctx1c

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

133 

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

135 return { 2 ctx1ac

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: 2 ctx1ba

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

145 

146 

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

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]): 2 ctx1ba

159 self.v = v 

160 

161 @overload 2 ctx1ba

162 def __getitem__(self, s: int) -> Q: ... 162 ↛ exitline 162 didn't return from function '__getitem__' because 2 ctx1ba

163 

164 @overload 2 ctx1ba

165 def __getitem__(self, s: slice) -> Sequence[Q]: ... 165 ↛ exitline 165 didn't return from function '__getitem__' because 2 ctx1ba

166 

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

168 return self.v[s] 

169 

170 def __len__(self) -> int: 2 ctx1ba

171 return len(self.v) 

172 

173 @classmethod 2 ctx1ba

174 def __get_pydantic_core_schema__( 2 ctx1ba

175 cls, source: Any, handler: GetCoreSchemaHandler 

176 ) -> core_schema.CoreSchema: 

177 args = get_args(source) 2 ctx1ad

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

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

180 

181 generic_type = args[0] 2 ctx1ad

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

183 scalar_schema = handler.generate_schema(generic_type) 2 ctx1ad

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

185 

186 return core_schema.no_info_after_validator_function( 2 ctx1ad

187 cls._scalar_to_sequence, handler(union_schema) 

188 ) 

189 

190 @classmethod 2 ctx1ba

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

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 true1 ctx1c

193 return v 1 ctx1c

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

195 

196 

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

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

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

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

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

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

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

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

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