Coverage for polar/kit/address.py: 66%

1from enum import StrEnum 

2from typing import TYPE_CHECKING, Annotated, Any, NotRequired, Self, TypedDict, cast 

3 

4import pycountry 

5from pydantic import BaseModel, BeforeValidator, Field, model_validator 

6from pydantic.json_schema import WithJsonSchema 

7from sqlalchemy.dialects.postgresql import JSONB 

8from sqlalchemy.engine.interfaces import Dialect 

9from sqlalchemy.types import TypeDecorator 

10 

11from polar.kit.schemas import EmptyStrToNone 

12 

13 

14class CountryData: 

15 alpha_2: str 

16 

17 

18_ALL_COUNTRIES: set[str] = { 

19 cast(CountryData, country).alpha_2 for country in pycountry.countries 

20} 

21_SUPPORTED_COUNTRIES: set[str] = _ALL_COUNTRIES - { 

22 # US Trade Embargos 

23 "CU", 

24 "IR", 

25 "KP", 

26 "SY", 

27 "RU", 

28} 

29ALL_COUNTRIES = sorted(_ALL_COUNTRIES) 

30SUPPORTED_COUNTRIES = sorted(_SUPPORTED_COUNTRIES) 

31 

32if TYPE_CHECKING: 32 ↛ 34line 32 didn't jump to line 34 because the condition on line 32 was never true

33 

34 class CountryAlpha2(StrEnum): 

35 pass 

36 

37 class CountryAlpha2Input(StrEnum): 

38 pass 

39else: 

40 CountryAlpha2 = Annotated[ 

41 StrEnum("CountryAlpha2", [(country, country) for country in ALL_COUNTRIES]), 

42 WithJsonSchema( 

43 { 

44 "type": "string", 

45 "title": "CountryAlpha2", 

46 "enum": ALL_COUNTRIES, 

47 "x-speakeasy-enums": ALL_COUNTRIES, 

48 } 

49 ), 

50 ] 

51 CountryAlpha2Input = Annotated[ 

52 StrEnum( 

53 "CountryAlpha2Input", 

54 [(country, country) for country in SUPPORTED_COUNTRIES], 

55 ), 

56 WithJsonSchema( 

57 { 

58 "type": "string", 

59 "title": "CountryAlpha2Input", 

60 "enum": SUPPORTED_COUNTRIES, 

61 "x-speakeasy-enums": SUPPORTED_COUNTRIES, 

62 } 

63 ), 

64 ] 

65 

66 

67class USState(StrEnum): 

68 US_AL = "US-AL" 

69 US_AK = "US-AK" 

70 US_AZ = "US-AZ" 

71 US_AR = "US-AR" 

72 US_CA = "US-CA" 

73 US_CO = "US-CO" 

74 US_CT = "US-CT" 

75 US_DE = "US-DE" 

76 US_FL = "US-FL" 

77 US_GA = "US-GA" 

78 US_HI = "US-HI" 

79 US_ID = "US-ID" 

80 US_IL = "US-IL" 

81 US_IN = "US-IN" 

82 US_IA = "US-IA" 

83 US_KS = "US-KS" 

84 US_KY = "US-KY" 

85 US_LA = "US-LA" 

86 US_ME = "US-ME" 

87 US_MD = "US-MD" 

88 US_MA = "US-MA" 

89 US_MI = "US-MI" 

90 US_MN = "US-MN" 

91 US_MS = "US-MS" 

92 US_MO = "US-MO" 

93 US_MT = "US-MT" 

94 US_NE = "US-NE" 

95 US_NV = "US-NV" 

96 US_NH = "US-NH" 

97 US_NJ = "US-NJ" 

98 US_NM = "US-NM" 

99 US_NY = "US-NY" 

100 US_NC = "US-NC" 

101 US_ND = "US-ND" 

102 US_OH = "US-OH" 

103 US_OK = "US-OK" 

104 US_OR = "US-OR" 

105 US_PA = "US-PA" 

106 US_RI = "US-RI" 

107 US_SC = "US-SC" 

108 US_SD = "US-SD" 

109 US_TN = "US-TN" 

110 US_TX = "US-TX" 

111 US_UT = "US-UT" 

112 US_VT = "US-VT" 

113 US_VA = "US-VA" 

114 US_WA = "US-WA" 

115 US_WV = "US-WV" 

116 US_WI = "US-WI" 

117 US_WY = "US-WY" 

118 US_DC = "US-DC" 

119 

120 

121class CAProvince(StrEnum): 

122 CA_AB = "CA-AB" 

123 CA_BC = "CA-BC" 

124 CA_MB = "CA-MB" 

125 CA_NB = "CA-NB" 

126 CA_NL = "CA-NL" 

127 CA_NS = "CA-NS" 

128 CA_ON = "CA-ON" 

129 CA_PE = "CA-PE" 

130 CA_QC = "CA-QC" 

131 CA_SK = "CA-SK" 

132 

133 

134class AddressDict(TypedDict): 

135 line1: NotRequired[str] 

136 line2: NotRequired[str] 

137 postal_code: NotRequired[str] 

138 city: NotRequired[str] 

139 state: NotRequired[str] 

140 country: str 

141 

142 

143class Address(BaseModel): 

144 line1: EmptyStrToNone | None = None 

145 line2: EmptyStrToNone | None = None 

146 postal_code: EmptyStrToNone | None = None 

147 city: EmptyStrToNone | None = None 

148 state: EmptyStrToNone | None = None 

149 country: CountryAlpha2 = Field(examples=["US", "SE", "FR"]) 

150 

151 @model_validator(mode="after") 

152 def validate_state(self) -> Self: 

153 if self.state is None: 

154 return self 

155 

156 # Normalize US and CA state with a prefix 

157 if self.country in {"US", "CA"}: 

158 if not self.state.startswith(f"{self.country}-"): 158 ↛ 159line 158 didn't jump to line 159 because the condition on line 158 was never true

159 self.state = f"{self.country}-{self.state}" 

160 

161 # Validate US and CA state 

162 if self.country == "US" and self.state not in USState: 162 ↛ 163line 162 didn't jump to line 163 because the condition on line 162 was never true

163 raise ValueError("Invalid US state") 

164 if self.country == "CA" and self.state not in CAProvince: 164 ↛ 165line 164 didn't jump to line 165 because the condition on line 164 was never true

165 raise ValueError("Invalid CA province") 

166 

167 return self 

168 

169 def to_dict(self) -> AddressDict: 

170 return cast(AddressDict, self.model_dump(exclude_none=True)) 

171 

172 def get_unprefixed_state(self) -> str | None: 

173 if self.state is None: 

174 return None 

175 if self.country in {"US", "CA"}: 

176 return self.state.split("-")[1] 

177 return self.state 

178 

179 def has_state(self) -> bool: 

180 return self.state is not None 

181 

182 def has_address(self) -> bool: 

183 return ( 

184 self.line1 is not None 

185 or self.line2 is not None 

186 or self.city is not None 

187 or self.postal_code is not None 

188 ) 

189 

190 def to_text(self) -> str: 

191 lines = [] 

192 if self.line1: 

193 lines.append(self.line1) 

194 if self.line2: 

195 lines.append(self.line2) 

196 

197 city_line = "" 

198 if self.city: 

199 city_line += self.city 

200 if self.state: 

201 state = pycountry.subdivisions.get(code=self.state) 

202 if state is not None: 

203 city_line += f", {cast(Any, state.name)}" 

204 else: 

205 city_line += f", {self.get_unprefixed_state()}" 

206 if self.postal_code: 

207 city_line += f" {self.postal_code}" 

208 if city_line: 

209 lines.append(city_line) 

210 

211 if self.country: 

212 country = pycountry.countries.get(alpha_2=self.country) 

213 if country is not None: 

214 lines.append(country.name) 

215 else: 

216 lines.append(self.country) 

217 

218 return "\n".join(lines) 

219 

220 

221class AddressInput(Address): 

222 country: Annotated[CountryAlpha2Input, BeforeValidator(str.upper)] = Field( # type: ignore 

223 examples=["US", "SE", "FR"] 

224 ) 

225 

226 

227class AddressType(TypeDecorator[Any]): 

228 impl = JSONB 

229 cache_ok = True 

230 

231 def process_bind_param(self, value: Any, dialect: Dialect) -> Any: 

232 if isinstance(value, Address): 232 ↛ 233line 232 didn't jump to line 233 because the condition on line 232 was never true

233 return value.model_dump(exclude_none=True) 

234 return value 

235 

236 def process_result_value(self, value: str | None, dialect: Dialect) -> Any: 

237 if value is not None: 

238 return Address.model_validate(value) 

239 return value