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

1from enum import StrEnum 2 ctx1ab

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

3 

4import pycountry 2 ctx1ab

5from pydantic import BaseModel, BeforeValidator, Field, model_validator 2 ctx1ab

6from pydantic.json_schema import WithJsonSchema 2 ctx1ab

7from sqlalchemy.dialects.postgresql import JSONB 2 ctx1ab

8from sqlalchemy.engine.interfaces import Dialect 2 ctx1ab

9from sqlalchemy.types import TypeDecorator 2 ctx1ab

10 

11from polar.kit.schemas import EmptyStrToNone 2 ctx1ab

12 

13 

14class CountryData: 2 ctx1ab

15 alpha_2: str 

16 

17 

18_ALL_COUNTRIES: set[str] = { 2 ctx1ab

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

20} 

21_SUPPORTED_COUNTRIES: set[str] = _ALL_COUNTRIES - { 2 ctx1ab

22 # US Trade Embargos 

23 "CU", 

24 "IR", 

25 "KP", 

26 "SY", 

27 "RU", 

28} 

29ALL_COUNTRIES = sorted(_ALL_COUNTRIES) 2 ctx1ab

30SUPPORTED_COUNTRIES = sorted(_SUPPORTED_COUNTRIES) 2 ctx1ab

31 

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

33 

34 class CountryAlpha2(StrEnum): 

35 pass 

36 

37 class CountryAlpha2Input(StrEnum): 

38 pass 

39else: 

40 CountryAlpha2 = Annotated[ 2 ctx1ab

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[ 2 ctx1ab

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): 2 ctx1ab

68 US_AL = "US-AL" 2 ctx1ab

69 US_AK = "US-AK" 2 ctx1ab

70 US_AZ = "US-AZ" 2 ctx1ab

71 US_AR = "US-AR" 2 ctx1ab

72 US_CA = "US-CA" 2 ctx1ab

73 US_CO = "US-CO" 2 ctx1ab

74 US_CT = "US-CT" 2 ctx1ab

75 US_DE = "US-DE" 2 ctx1ab

76 US_FL = "US-FL" 2 ctx1ab

77 US_GA = "US-GA" 2 ctx1ab

78 US_HI = "US-HI" 2 ctx1ab

79 US_ID = "US-ID" 2 ctx1ab

80 US_IL = "US-IL" 2 ctx1ab

81 US_IN = "US-IN" 2 ctx1ab

82 US_IA = "US-IA" 2 ctx1ab

83 US_KS = "US-KS" 2 ctx1ab

84 US_KY = "US-KY" 2 ctx1ab

85 US_LA = "US-LA" 2 ctx1ab

86 US_ME = "US-ME" 2 ctx1ab

87 US_MD = "US-MD" 2 ctx1ab

88 US_MA = "US-MA" 2 ctx1ab

89 US_MI = "US-MI" 2 ctx1ab

90 US_MN = "US-MN" 2 ctx1ab

91 US_MS = "US-MS" 2 ctx1ab

92 US_MO = "US-MO" 2 ctx1ab

93 US_MT = "US-MT" 2 ctx1ab

94 US_NE = "US-NE" 2 ctx1ab

95 US_NV = "US-NV" 2 ctx1ab

96 US_NH = "US-NH" 2 ctx1ab

97 US_NJ = "US-NJ" 2 ctx1ab

98 US_NM = "US-NM" 2 ctx1ab

99 US_NY = "US-NY" 2 ctx1ab

100 US_NC = "US-NC" 2 ctx1ab

101 US_ND = "US-ND" 2 ctx1ab

102 US_OH = "US-OH" 2 ctx1ab

103 US_OK = "US-OK" 2 ctx1ab

104 US_OR = "US-OR" 2 ctx1ab

105 US_PA = "US-PA" 2 ctx1ab

106 US_RI = "US-RI" 2 ctx1ab

107 US_SC = "US-SC" 2 ctx1ab

108 US_SD = "US-SD" 2 ctx1ab

109 US_TN = "US-TN" 2 ctx1ab

110 US_TX = "US-TX" 2 ctx1ab

111 US_UT = "US-UT" 2 ctx1ab

112 US_VT = "US-VT" 2 ctx1ab

113 US_VA = "US-VA" 2 ctx1ab

114 US_WA = "US-WA" 2 ctx1ab

115 US_WV = "US-WV" 2 ctx1ab

116 US_WI = "US-WI" 2 ctx1ab

117 US_WY = "US-WY" 2 ctx1ab

118 US_DC = "US-DC" 2 ctx1ab

119 

120 

121class CAProvince(StrEnum): 2 ctx1ab

122 CA_AB = "CA-AB" 2 ctx1ab

123 CA_BC = "CA-BC" 2 ctx1ab

124 CA_MB = "CA-MB" 2 ctx1ab

125 CA_NB = "CA-NB" 2 ctx1ab

126 CA_NL = "CA-NL" 2 ctx1ab

127 CA_NS = "CA-NS" 2 ctx1ab

128 CA_ON = "CA-ON" 2 ctx1ab

129 CA_PE = "CA-PE" 2 ctx1ab

130 CA_QC = "CA-QC" 2 ctx1ab

131 CA_SK = "CA-SK" 2 ctx1ab

132 

133 

134class AddressDict(TypedDict): 2 ctx1ab

135 line1: NotRequired[str] 2 ctx1ab

136 line2: NotRequired[str] 2 ctx1ab

137 postal_code: NotRequired[str] 2 ctx1ab

138 city: NotRequired[str] 2 ctx1ab

139 state: NotRequired[str] 2 ctx1ab

140 country: str 2 ctx1ab

141 

142 

143class Address(BaseModel): 2 ctx1ab

144 line1: EmptyStrToNone | None = None 2 ctx1ab

145 line2: EmptyStrToNone | None = None 2 ctx1ab

146 postal_code: EmptyStrToNone | None = None 2 ctx1ab

147 city: EmptyStrToNone | None = None 2 ctx1ab

148 state: EmptyStrToNone | None = None 2 ctx1ab

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

150 

151 @model_validator(mode="after") 2 ctx1ab

152 def validate_state(self) -> Self: 2 ctx1ab

153 if self.state is None: 3 ctx1abc

154 return self 1 ctx1c

155 

156 # Normalize US and CA state with a prefix 

157 if self.country in {"US", "CA"}: 157 ↛ 162line 157 didn't jump to line 162 because the condition on line 157 was always true2 ctx1ab

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 true2 ctx1ab

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 true2 ctx1ab

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 true2 ctx1ab

165 raise ValueError("Invalid CA province") 

166 

167 return self 2 ctx1ab

168 

169 def to_dict(self) -> AddressDict: 2 ctx1ab

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

171 

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

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

180 return self.state is not None 

181 

182 def has_address(self) -> bool: 2 ctx1ab

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

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): 2 ctx1ab

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

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

224 ) 

225 

226 

227class AddressType(TypeDecorator[Any]): 2 ctx1ab

228 impl = JSONB 2 ctx1ab

229 cache_ok = True 2 ctx1ab

230 

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

232 if isinstance(value, Address): 

233 return value.model_dump(exclude_none=True) 

234 return value 

235 

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

237 if value is not None: 

238 return Address.model_validate(value) 

239 return value