Coverage for src/usaspending/queries/filters.py: 96%

1from __future__ import annotations 

2 

3import datetime 

4from abc import ABC, abstractmethod 

5from dataclasses import dataclass, field 

6from enum import Enum 

7from typing import Any, ClassVar, Literal, Optional 

8 

9# ============================================================================== 

10# Award Type Code Constants 

11# ============================================================================== 

12 

13# Contract award type codes 

14CONTRACT_CODES = frozenset({"A", "B", "C", "D"}) 

15 

16# IDV award type codes 

17IDV_CODES = frozenset( 

18 {"IDV_A", "IDV_B", "IDV_B_A", "IDV_B_B", "IDV_B_C", "IDV_C", "IDV_D", "IDV_E"} 

19) 

20 

21# Loan award type codes 

22LOAN_CODES = frozenset({"07", "08"}) 

23 

24# Grant award type codes 

25GRANT_CODES = frozenset({"02", "03", "04", "05"}) 

26 

27# Direct payment award type codes 

28DIRECT_PAYMENT_CODES = frozenset({"06", "10"}) 

29 

30# Other award type codes that do not fit into the above categories 

31OTHER_CODES = frozenset({"09", "11", "-1"}) 

32 

33# All valid award type codes 

34ALL_AWARD_CODES = CONTRACT_CODES | IDV_CODES | LOAN_CODES | GRANT_CODES 

35 

36# ============================================================================== 

37# Helper Enums and Dataclasses 

38# ============================================================================== 

39 

40 

41class AgencyType(Enum): 

42 """Enumeration for agency types.""" 

43 

44 AWARDING = "awarding" 

45 FUNDING = "funding" 

46 

47 

48class AgencyTier(Enum): 

49 """Enumeration for agency tiers.""" 

50 

51 TOPTIER = "toptier" 

52 SUBTIER = "subtier" 

53 

54 

55class LocationScope(Enum): 

56 """Enumeration for location scopes.""" 

57 

58 DOMESTIC = "domestic" 

59 FOREIGN = "foreign" 

60 

61 

62class AwardDateType(Enum): 

63 """Enumeration for award search date types.""" 

64 

65 ACTION_DATE = "action_date" 

66 DATE_SIGNED = "date_signed" 

67 LAST_MODIFIED = "last_modified_date" 

68 NEW_AWARDS_ONLY = "new_awards_only" 

69 

70 

71@dataclass(frozen=True) 

72class LocationSpec: 

73 """Represents a standard location specification for Place of Performance or Recipient filters.""" 

74 

75 country_code: str 

76 state_code: Optional[str] = None 

77 county_code: Optional[str] = None 

78 city_name: Optional[str] = None 

79 district_original: Optional[str] = None 

80 district_current: Optional[str] = None 

81 zip_code: Optional[str] = None 

82 

83 def to_dict(self) -> dict[str, str]: 

84 """Serializes the location to the dictionary format required by the API.""" 

85 data = {"country": self.country_code} 

86 if self.state_code: 

87 data["state"] = self.state_code 

88 if self.county_code: 

89 data["county"] = self.county_code 

90 if self.city_name: 

91 data["city"] = self.city_name 

92 if self.district_original: 

93 data["district_original"] = self.district_original 

94 if self.district_current: 

95 data["district_current"] = self.district_current 

96 if self.zip_code: 

97 data["zip"] = self.zip_code 

98 return data 

99 

100 

101# ============================================================================== 

102# Base Filter Abstraction 

103# ============================================================================== 

104 

105 

106class BaseFilter(ABC): 

107 """Abstract base class for all query filter types.""" 

108 

109 key: ClassVar[str] 

110 

111 @abstractmethod 

112 def to_dict(self) -> dict[str, Any]: 

113 """Converts the filter to its dictionary representation for the API.""" 

114 pass 

115 

116 

117# ============================================================================== 

118# Individual Filter Implementations 

119# ============================================================================== 

120 

121 

122@dataclass(frozen=True) 

123class KeywordsFilter(BaseFilter): 

124 """Filter by a list of keywords.""" 

125 

126 key: ClassVar[str] = "keywords" 

127 values: list[str] 

128 

129 def to_dict(self) -> dict[str, list[str]]: 

130 return {self.key: self.values} 

131 

132 

133@dataclass(frozen=True) 

134class TimePeriodFilter(BaseFilter): 

135 """Filter by a date range.""" 

136 

137 key: ClassVar[str] = "time_period" 

138 start_date: datetime.date 

139 end_date: datetime.date 

140 date_type: Optional[AwardDateType] = None 

141 

142 def to_dict(self) -> dict[str, list[dict[str, str]]]: 

143 period: dict[str, str] = { 

144 "start_date": self.start_date.strftime("%Y-%m-%d"), 

145 "end_date": self.end_date.strftime("%Y-%m-%d"), 

146 } 

147 if self.date_type: 

148 period["date_type"] = self.date_type.value 

149 return {self.key: [period]} 

150 

151 

152@dataclass(frozen=True) 

153class LocationScopeFilter(BaseFilter): 

154 """Filter by domestic or foreign scope for location.""" 

155 

156 key: Literal["place_of_performance_scope", "recipient_scope"] 

157 scope: LocationScope 

158 

159 def to_dict(self) -> dict[str, str]: 

160 return {self.key: self.scope.value} 

161 

162 

163@dataclass(frozen=True) 

164class LocationFilter(BaseFilter): 

165 """Filter by one or more specific geographic locations.""" 

166 

167 key: Literal["place_of_performance_locations", "recipient_locations"] 

168 locations: list[LocationSpec] 

169 

170 def to_dict(self) -> dict[str, list[dict[str, str]]]: 

171 return {self.key: [loc.to_dict() for loc in self.locations]} 

172 

173 

174@dataclass(frozen=True) 

175class AgencyFilter(BaseFilter): 

176 """Filter by an awarding or funding agency.""" 

177 

178 key: ClassVar[str] = "agencies" 

179 agency_type: AgencyType 

180 tier: AgencyTier 

181 name: str 

182 

183 def to_dict(self) -> dict[str, list[dict[str, str]]]: 

184 agency_object = { 

185 "type": self.agency_type.value, 

186 "tier": self.tier.value, 

187 "name": self.name, 

188 } 

189 return {self.key: [agency_object]} 

190 

191 

192@dataclass(frozen=True) 

193class SimpleListFilter(BaseFilter): 

194 """A generic filter for API keys that accept a list of string values.""" 

195 

196 key: str 

197 values: list[str] 

198 

199 def to_dict(self) -> dict[str, list[str]]: 

200 return {self.key: self.values} 

201 

202 

203@dataclass(frozen=True) 

204class AwardAmount: 

205 """Represents a single award amount range for filtering.""" 

206 

207 lower_bound: Optional[float] = None 

208 upper_bound: Optional[float] = None 

209 

210 def to_dict(self) -> dict[str, float]: 

211 data = {} 

212 if self.lower_bound is not None: 

213 data["lower_bound"] = self.lower_bound 

214 if self.upper_bound is not None: 

215 data["upper_bound"] = self.upper_bound 

216 return data 

217 

218 

219@dataclass(frozen=True) 

220class AwardAmountFilter(BaseFilter): 

221 """Filter by one or more award amount ranges.""" 

222 

223 key: ClassVar[str] = "award_amounts" 

224 amounts: list[AwardAmount] 

225 

226 def to_dict(self) -> dict[str, list[dict[str, float]]]: 

227 return {self.key: [amount.to_dict() for amount in self.amounts]} 

228 

229 

230@dataclass(frozen=True) 

231class TieredCodeFilter(BaseFilter): 

232 """Handles filters with a 'require' and 'exclude' structure like NAICS.""" 

233 

234 key: Literal["naics_codes", "psc_codes", "tas_codes"] 

235 require: list[list[str]] = field(default_factory=list) 

236 exclude: list[list[str]] = field(default_factory=list) 

237 

238 def to_dict(self) -> dict[str, dict[str, list[list[str]]]]: 

239 data = {} 

240 if self.require: 

241 data["require"] = self.require 

242 if self.exclude: 

243 data["exclude"] = self.exclude 

244 return {self.key: data} 

245 

246 

247@dataclass(frozen=True) 

248class TreasuryAccountComponentsFilter(BaseFilter): 

249 """Filter by specific components of a Treasury Account.""" 

250 

251 key: ClassVar[str] = "treasury_account_components" 

252 components: list[dict[str, str]] 

253 

254 def to_dict(self) -> dict[str, list[dict[str, str]]]: 

255 return {self.key: self.components}