Coverage for src/usaspending/queries/awards_search.py: 99%
221 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-09-03 17:15 -0700
1from __future__ import annotations
3import datetime
4from typing import Any, Optional, Union
6from usaspending.client import USASpending
7from usaspending.exceptions import ValidationError
8from usaspending.models.award_factory import create_award
9from usaspending.models import Award
10from usaspending.models.contract import Contract
11from usaspending.models.grant import Grant
12from usaspending.models.idv import IDV
13from usaspending.models.loan import Loan
14from usaspending.queries.query_builder import QueryBuilder
15from usaspending.logging_config import USASpendingLogger
16from usaspending.queries.filters import (
17 AgencyFilter,
18 AgencyTier,
19 AgencyType,
20 AwardAmount,
21 AwardAmountFilter,
22 AwardDateType,
23 KeywordsFilter,
24 LocationSpec,
25 LocationFilter,
26 LocationScope,
27 LocationScopeFilter,
28 SimpleListFilter,
29 TieredCodeFilter,
30 TimePeriodFilter,
31 TreasuryAccountComponentsFilter,
32)
34# Import award type codes from config
35# These are defined by USASpending.gov and represent different categories of awards
36from ..config import (
37 CONTRACT_CODES,
38 IDV_CODES,
39 LOAN_CODES,
40 GRANT_CODES,
41 DIRECT_PAYMENT_CODES,
42 OTHER_CODES,
43 AWARD_TYPE_GROUPS,
44)
46logger = USASpendingLogger.get_logger(__name__)
49class AwardsSearch(QueryBuilder["Award"]):
50 """
51 Builds and executes a spending_by_award search query, allowing for complex
52 filtering on award data. This class follows a fluent interface pattern.
53 """
55 def __init__(self, client: USASpending):
56 """
57 Initializes the AwardsSearch query builder.
59 Args:
60 client: The USASpending client instance.
61 """
62 super().__init__(client)
64 @property
65 def _endpoint(self) -> str:
66 """The API endpoint for this query."""
67 return "/search/spending_by_award/"
69 def _clone(self) -> AwardsSearch:
70 """Creates an immutable copy of the query builder."""
71 clone = super()._clone()
72 clone._filter_objects = self._filter_objects.copy()
73 return clone
75 def _build_payload(self, page: int) -> dict[str, Any]:
76 """Constructs the final API request payload from the filter objects."""
78 final_filters = self._aggregate_filters()
80 # The 'award_type_codes' filter is required by the API.
81 if "award_type_codes" not in final_filters:
82 raise ValidationError(
83 "A filter for 'award_type_codes' is required. "
84 "Use the .with_award_types() method."
85 )
87 payload = {
88 "filters": final_filters,
89 "fields": self._get_fields(),
90 "limit": self._get_effective_page_size(),
91 "page": page,
92 }
93 return payload
95 def _transform_result(self, result: dict[str, Any]) -> Award:
96 """Transforms a single API result item into an Award model."""
97 # Get award type codes from current filters
98 award_type_codes = self._get_award_type_codes()
100 # If we're filtering for a single award type category, add it to the result
101 # This ensures the correct Award subclass is created even when the API
102 # response doesn't include explicit type information
103 if award_type_codes:
104 if award_type_codes.issubset(CONTRACT_CODES):
105 result["category"] = "contract"
106 elif award_type_codes.issubset(IDV_CODES):
107 result["category"] = "idv"
108 elif award_type_codes.issubset(GRANT_CODES):
109 result["category"] = "grant"
110 elif award_type_codes.issubset(LOAN_CODES):
111 result["category"] = "loan"
113 return create_award(result, self._client)
115 def _get_award_type_codes(self) -> set[str]:
116 """Extract award type codes from current filters."""
117 for filter_obj in self._filter_objects:
118 filter_dict = filter_obj.to_dict()
119 if "award_type_codes" in filter_dict:
120 return set(filter_dict["award_type_codes"])
121 return set()
123 def _validate_single_award_type_category(self, new_codes: set[str]) -> None:
124 """
125 Validate that only one category of award types is present.
127 Args:
128 new_codes: New award type codes being added
130 Raises:
131 ValidationError: If mixing award type categories
132 """
133 existing_codes = self._get_award_type_codes()
134 all_codes = existing_codes | new_codes
136 if not all_codes:
137 return
139 # Check how many categories are represented using the config mapping
140 categories_present = 0
141 category_names = []
143 for category_name, codes in AWARD_TYPE_GROUPS.items():
144 if all_codes & frozenset(codes.keys()):
145 categories_present += 1
146 category_names.append(category_name)
148 if categories_present > 1:
149 raise ValidationError(
150 f"Cannot mix different award type categories: {', '.join(category_names)}. "
151 "Use separate queries for each award type category."
152 )
154 def count(self) -> int:
155 """
156 Get the total count of results without fetching all items.
158 Returns:
159 The total number of matching awards for the selected award type category.
160 """
161 logger.debug(f"{self.__class__.__name__}.count() called")
163 # Aggregate filters to prepare for the count request
164 final_filters = self._aggregate_filters()
166 # The 'award_type_codes' filter is required by the API.
167 if "award_type_codes" not in final_filters:
168 raise ValidationError(
169 "A filter for 'award_type_codes' is required. "
170 "Use the .with_award_types() method."
171 )
173 # Make the API call to count awards by type
174 results = self.count_awards_by_type()
176 # Get the award type codes to determine which category to count
177 award_type_codes = self._get_award_type_codes()
179 # Determine the category based on award type codes
180 category = self._get_award_type_category(award_type_codes)
182 # Extract the count for the specific category
183 total = results.get(category, 0)
185 logger.info(f"{self.__class__.__name__}.count() = {total} ({category})")
186 return total
188 def count_awards_by_type(self) -> dict[str, int]:
189 """ Shared logic that calls the awards count endpoint.
191 Returns:
192 A dictionary mapping award type categories to their result counts
193 for the matching filter set.
194 """
195 endpoint = "/search/spending_by_award_count/"
196 final_filters = self._aggregate_filters()
198 payload = {
199 "filters": final_filters,
200 }
202 from ..logging_config import log_query_execution
204 log_query_execution(
205 logger, "AwardsSearch._count_awards_by_type", len(self._filter_objects), endpoint
206 )
208 # Send the request to the count endpoint
209 response = self._client._make_request("POST", endpoint, json=payload)
210 results = response.get("results", {})
212 return results
214 def _get_award_type_category(self, award_type_codes: set[str]) -> str:
215 """
216 Determine the award type category based on the award type codes.
218 Args:
219 award_type_codes: Set of award type codes
221 Returns:
222 The category name as used in the count endpoint response
223 """
224 # Map config category names to API response names
225 category_mapping = {
226 "contracts": "contracts",
227 "idvs": "idvs",
228 "loans": "loans",
229 "grants": "grants",
230 "direct_payments": "direct_payments",
231 "other_assistance": "other",
232 }
234 for category_name, codes in AWARD_TYPE_GROUPS.items():
235 if award_type_codes & frozenset(codes.keys()):
236 return category_mapping[category_name]
238 # Fail hard if no valid award type category is found
239 raise ValidationError("No valid award type category found. ")
241 def _get_fields(self) -> list[str]:
242 """
243 Determines the list of fields to request based on award type filters.
245 Returns different field sets depending on the award type codes:
246 - Contracts (A, B, C, D): Include contract-specific fields
247 - IDV (IDV_A, IDV_B, etc.): Include IDV-specific fields
248 - Loans (07, 08): Include loan-specific fields
249 - Grants/Assistance (02, 03, 04, 05, 06, 09, 10, 11, -1): Include assistance fields
250 """
251 # Start with base fields from Award model
252 base_fields = Award.SEARCH_FIELDS.copy()
254 # Get award type codes from filters
255 award_types = self._get_award_type_codes()
256 additional_fields = []
258 # Check each category and add appropriate fields based on model
259 for category_name, codes in AWARD_TYPE_GROUPS.items():
260 if award_types & frozenset(codes.keys()):
261 if category_name == "contracts":
262 # Use Contract.SEARCH_FIELDS but exclude base fields
263 additional_fields.extend(
264 [f for f in Contract.SEARCH_FIELDS if f not in base_fields]
265 )
266 elif category_name == "idvs":
267 # Use IDV.SEARCH_FIELDS but exclude base fields
268 additional_fields.extend(
269 [f for f in IDV.SEARCH_FIELDS if f not in base_fields]
270 )
271 elif category_name == "loans":
272 # Use Loan.SEARCH_FIELDS but exclude base fields
273 additional_fields.extend(
274 [f for f in Loan.SEARCH_FIELDS if f not in base_fields]
275 )
276 elif category_name in ["grants", "direct_payments", "other_assistance"]:
277 # Use Grant.SEARCH_FIELDS but exclude base fields
278 additional_fields.extend(
279 [f for f in Grant.SEARCH_FIELDS if f not in base_fields]
280 )
282 # Combine base fields with additional fields, removing duplicates
283 all_fields = base_fields + additional_fields
284 return list(
285 dict.fromkeys(all_fields)
286 ) # Remove duplicates while preserving order
288 # ==========================================================================
289 # Filter Methods
290 # ==========================================================================
292 def with_keywords(self, *keywords: str) -> AwardsSearch:
293 """
294 Filter by a list of keywords.
296 Args:
297 *keywords: One or more keywords to search for.
299 Returns:
300 A new `AwardsSearch` instance with the filter applied.
301 """
302 clone = self._clone()
303 clone._filter_objects.append(KeywordsFilter(values=list(keywords)))
304 return clone
306 def in_time_period(
307 self,
308 start_date: Union[datetime.date, str],
309 end_date: Union[datetime.date, str],
310 new_awards_only: bool = False,
311 date_type: Optional[AwardDateType] = None,
312 ) -> AwardsSearch:
313 """
314 Filter by a specific date range.
316 Args:
317 start_date: The start date of the period (datetime.date or string in "YYYY-MM-DD" format).
318 end_date: The end date of the period (datetime.date or string in "YYYY-MM-DD" format).
319 new_awards_only: If True, filters by awards with a start date within the given range.
320 date_type: The type of date to filter on (e.g., action_date).
322 Returns:
323 A new `AwardsSearch` instance with the filter applied.
325 Raises:
326 ValidationError: If string dates are not in valid "YYYY-MM-DD" format.
327 """
329 # Parse string dates if needed
330 if isinstance(start_date, str):
331 try:
332 start_date = datetime.datetime.strptime(start_date, "%Y-%m-%d").date()
333 except ValueError:
334 raise ValidationError(
335 f"Invalid start_date format: '{start_date}'. Expected 'YYYY-MM-DD'."
336 )
338 if isinstance(end_date, str):
339 try:
340 end_date = datetime.datetime.strptime(end_date, "%Y-%m-%d").date()
341 except ValueError:
342 raise ValidationError(
343 f"Invalid end_date format: '{end_date}'. Expected 'YYYY-MM-DD'."
344 )
346 # If convenience flag is set, use NEW_AWARDS_ONLY date type
347 # and override any provided date_type
348 if new_awards_only:
349 date_type = AwardDateType.NEW_AWARDS_ONLY
350 clone = self._clone()
351 clone._filter_objects.append(
352 TimePeriodFilter(
353 start_date=start_date, end_date=end_date, date_type=date_type
354 )
355 )
356 return clone
358 def for_fiscal_year(
359 self,
360 year: int,
361 new_awards_only: bool = False,
362 date_type: Optional[AwardDateType] = None,
363 ) -> AwardsSearch:
364 """
365 Adds a time period filter for a single US government fiscal year
366 (October 1 to September 30).
368 Args:
369 year: The fiscal year to filter by.
370 new_awards_only: If True, filters by awards with a start date within the FY
371 date_type: The type of date to filter on (e.g., action_date).
373 Returns:
374 A new `AwardsSearch` instance with the fiscal year filter applied.
375 """
376 start_date = datetime.date(year - 1, 10, 1)
377 end_date = datetime.date(year, 9, 30)
378 return self.in_time_period(
379 start_date=start_date,
380 end_date=end_date,
381 new_awards_only=new_awards_only,
382 date_type=date_type,
383 )
385 def with_place_of_performance_scope(self, scope: LocationScope) -> AwardsSearch:
386 """
387 Filter awards by domestic or foreign place of performance.
389 Args:
390 scope: The scope, either DOMESTIC or FOREIGN.
392 Returns:
393 A new `AwardsSearch` instance with the filter applied.
394 """
395 clone = self._clone()
396 clone._filter_objects.append(
397 LocationScopeFilter(key="place_of_performance_scope", scope=scope)
398 )
399 return clone
401 def with_place_of_performance_locations(self, *locations: LocationSpec) -> AwardsSearch:
402 """
403 Filter by one or more specific geographic places of performance.
405 Args:
406 *locations: One or more `LocationSpec` objects.
408 Returns:
409 A new `AwardsSearch` instance with the filter applied.
410 """
411 clone = self._clone()
412 clone._filter_objects.append(
413 LocationFilter(
414 key="place_of_performance_locations", locations=list(locations)
415 )
416 )
417 return clone
419 def for_agency(
420 self,
421 name: str,
422 agency_type: AgencyType = AgencyType.AWARDING,
423 tier: AgencyTier = AgencyTier.TOPTIER,
424 ) -> AwardsSearch:
425 """
426 Filter by a specific awarding or funding agency.
428 Args:
429 name: The name of the agency.
430 agency_type: The type of agency (AWARDING or FUNDING).
431 tier: The agency tier (TOPTIER or SUBTIER).
433 Returns:
434 A new `AwardsSearch` instance with the filter applied.
435 """
436 clone = self._clone()
437 clone._filter_objects.append(
438 AgencyFilter(agency_type=agency_type, tier=tier, name=name)
439 )
440 return clone
442 def with_recipient_search_text(self, *search_terms: str) -> AwardsSearch:
443 """
444 Filter by recipient name, UEI, or DUNS.
446 Args:
447 *search_terms: Text to search for across recipient identifiers.
449 Returns:
450 A new `AwardsSearch` instance with the filter applied.
451 """
452 clone = self._clone()
453 clone._filter_objects.append(
454 SimpleListFilter(key="recipient_search_text", values=list(search_terms))
455 )
456 return clone
458 def with_recipient_scope(self, scope: LocationScope) -> AwardsSearch:
459 """
460 Filter recipients by domestic or foreign scope.
462 Args:
463 scope: The scope, either DOMESTIC or FOREIGN.
465 Returns:
466 A new `AwardsSearch` instance with the filter applied.
467 """
468 clone = self._clone()
469 clone._filter_objects.append(
470 LocationScopeFilter(key="recipient_scope", scope=scope)
471 )
472 return clone
474 def with_recipient_locations(self, *locations: LocationSpec) -> AwardsSearch:
475 """
476 Filter by one or more specific recipient locations.
478 Args:
479 *locations: One or more `LocationSpec` objects.
481 Returns:
482 A new `AwardsSearch` instance with the filter applied.
483 """
484 clone = self._clone()
485 clone._filter_objects.append(
486 LocationFilter(key="recipient_locations", locations=list(locations))
487 )
488 return clone
490 def with_recipient_types(self, *type_names: str) -> AwardsSearch:
491 """
492 Filter by one or more recipient or business types.
494 Args:
495 *type_names: The names of the recipient types (e.g., "small_business").
497 Returns:
498 A new `AwardsSearch` instance with the filter applied.
499 """
500 clone = self._clone()
501 clone._filter_objects.append(
502 SimpleListFilter(key="recipient_type_names", values=list(type_names))
503 )
504 return clone
506 def with_award_types(self, *award_codes: str) -> AwardsSearch:
507 """
508 Filter by one or more award type codes. This filter is **required**.
510 Args:
511 *award_codes: A sequence of award type codes (e.g., "A", "B", "02").
513 Returns:
514 A new `AwardsSearch` instance with the filter applied.
516 Raises:
517 ValidationError: If mixing different award type categories.
518 """
519 new_codes = set(award_codes)
520 self._validate_single_award_type_category(new_codes)
522 clone = self._clone()
523 clone._filter_objects.append(
524 SimpleListFilter(key="award_type_codes", values=list(award_codes))
525 )
526 return clone
528 def contracts(self) -> AwardsSearch:
529 """
530 Filter to search for contract awards only (types A, B, C, D).
532 Returns:
533 A new `AwardsSearch` instance configured for contract awards.
534 """
535 return self.with_award_types(*CONTRACT_CODES)
537 def idvs(self) -> AwardsSearch:
538 """
539 Filter to search for IDV awards only (types IDV_A, IDV_B, etc.).
541 Returns:
542 A new `AwardsSearch` instance configured for IDV awards.
543 """
544 return self.with_award_types(*IDV_CODES)
546 def loans(self) -> AwardsSearch:
547 """
548 Filter to search for loan awards only (types 07, 08).
550 Returns:
551 A new `AwardsSearch` instance configured for loan awards.
552 """
553 return self.with_award_types(*LOAN_CODES)
555 def grants(self) -> AwardsSearch:
556 """
557 Filter to search for grant and assistance awards only (types 02, 03, 04, 05).
559 Returns:
560 A new `AwardsSearch` instance configured for grant/assistance awards.
561 """
562 return self.with_award_types(*GRANT_CODES)
564 def direct_payments(self) -> AwardsSearch:
565 """
566 Filter to search for direct payment awards only (types 06, 10).
568 Returns:
569 A new `AwardsSearch` instance configured for direct payment awards.
570 """
571 return self.with_award_types(*DIRECT_PAYMENT_CODES)
573 def other(self) -> AwardsSearch:
574 """
575 Filter to search for other assistance awards only (types 09, 11, -1).
577 Returns:
578 A new `AwardsSearch` instance configured for other assistance awards.
579 """
580 return self.with_award_types(*OTHER_CODES)
582 def with_award_ids(self, *award_ids: str) -> AwardsSearch:
583 """
584 Filter by specific award IDs (FAIN, PIID, URI).
586 Args:
587 *award_ids: The exact award IDs to search for.
589 Returns:
590 A new `AwardsSearch` instance with the filter applied.
591 """
592 clone = self._clone()
593 clone._filter_objects.append(
594 SimpleListFilter(key="award_ids", values=list(award_ids))
595 )
596 return clone
598 def with_award_amounts(self, *amounts: AwardAmount) -> AwardsSearch:
599 """
600 Filter by one or more award amount ranges.
602 Args:
603 *amounts: One or more `AwardAmount` objects defining the ranges.
605 Returns:
606 A new `AwardsSearch` instance with the filter applied.
607 """
608 clone = self._clone()
609 clone._filter_objects.append(AwardAmountFilter(amounts=list(amounts)))
610 return clone
612 def with_cfda_numbers(self, *program_numbers: str) -> AwardsSearch:
613 """
614 Filter by one or more CFDA program numbers.
616 Args:
617 *program_numbers: The CFDA numbers to filter by.
619 Returns:
620 A new `AwardsSearch` instance with the filter applied.
621 """
622 clone = self._clone()
623 clone._filter_objects.append(
624 SimpleListFilter(key="program_numbers", values=list(program_numbers))
625 )
626 return clone
628 def with_naics_codes(
629 self,
630 require: Optional[list[str]] = None,
631 exclude: Optional[list[str]] = None,
632 ) -> AwardsSearch:
633 """
634 Filter by NAICS codes, including or excluding specific codes.
636 Args:
637 require: A list of NAICS codes to require.
638 exclude: A list of NAICS codes to exclude.
640 Returns:
641 A new `AwardsSearch` instance with the filter applied.
642 """
643 clone = self._clone()
644 # The API expects a list of lists, but for NAICS, each list contains one element.
645 require_list = [[code] for code in require] if require else []
646 exclude_list = [[code] for code in exclude] if exclude else []
647 clone._filter_objects.append(
648 TieredCodeFilter(
649 key="naics_codes", require=require_list, exclude=exclude_list
650 )
651 )
652 return clone
654 def with_psc_codes(
655 self,
656 require: Optional[list[list[str]]] = None,
657 exclude: Optional[list[list[str]]] = None,
658 ) -> AwardsSearch:
659 """
660 Filter by Product and Service Codes (PSC), including or excluding codes.
662 Args:
663 require: A list of PSC code paths to require.
664 exclude: A list of PSC code paths to exclude.
666 Returns:
667 A new `AwardsSearch` instance with the filter applied.
668 """
669 clone = self._clone()
670 clone._filter_objects.append(
671 TieredCodeFilter(
672 key="psc_codes",
673 require=require or [],
674 exclude=exclude or [],
675 )
676 )
677 return clone
679 def with_contract_pricing_types(self, *type_codes: str) -> AwardsSearch:
680 """
681 Filter by one or more contract pricing type codes.
683 Args:
684 *type_codes: The contract pricing type codes.
686 Returns:
687 A new `AwardsSearch` instance with the filter applied.
688 """
689 clone = self._clone()
690 clone._filter_objects.append(
691 SimpleListFilter(key="contract_pricing_type_codes", values=list(type_codes))
692 )
693 return clone
695 def with_set_aside_types(self, *type_codes: str) -> AwardsSearch:
696 """
697 Filter by one or more set-aside type codes.
699 Args:
700 *type_codes: The set-aside type codes.
702 Returns:
703 A new `AwardsSearch` instance with the filter applied.
704 """
705 clone = self._clone()
706 clone._filter_objects.append(
707 SimpleListFilter(key="set_aside_type_codes", values=list(type_codes))
708 )
709 return clone
711 def with_extent_competed_types(self, *type_codes: str) -> AwardsSearch:
712 """
713 Filter by one or more extent competed type codes.
715 Args:
716 *type_codes: The extent competed type codes.
718 Returns:
719 A new `AwardsSearch` instance with the filter applied.
720 """
721 clone = self._clone()
722 clone._filter_objects.append(
723 SimpleListFilter(key="extent_competed_type_codes", values=list(type_codes))
724 )
725 return clone
727 def with_tas_codes(
728 self,
729 require: Optional[list[list[str]]] = None,
730 exclude: Optional[list[list[str]]] = None,
731 ) -> AwardsSearch:
732 """
733 Filter by Treasury Account Symbols (TAS), including or excluding codes.
735 Args:
736 require: A list of TAS code paths to require.
737 exclude: A list of TAS code paths to exclude.
739 Returns:
740 A new `AwardsSearch` instance with the filter applied.
741 """
742 clone = self._clone()
743 clone._filter_objects.append(
744 TieredCodeFilter(
745 key="tas_codes",
746 require=require or [],
747 exclude=exclude or [],
748 )
749 )
750 return clone
752 def with_treasury_account_components(
753 self, *components: dict[str, str]
754 ) -> AwardsSearch:
755 """
756 Filter by specific components of a Treasury Account.
758 Args:
759 *components: Dictionaries representing TAS components (aid, main, etc.).
761 Returns:
762 A new `AwardsSearch` instance with the filter applied.
763 """
764 clone = self._clone()
765 clone._filter_objects.append(
766 TreasuryAccountComponentsFilter(components=list(components))
767 )
768 return clone
770 def with_def_codes(self, *def_codes: str) -> AwardsSearch:
771 """
772 Filter by one or more Disaster Emergency Fund (DEF) codes.
774 Args:
775 *def_codes: The DEF codes (e.g., "L", "M", "N").
777 Returns:
778 A new `AwardsSearch` instance with the filter applied.
779 """
780 clone = self._clone()
781 clone._filter_objects.append(
782 SimpleListFilter(key="def_codes", values=list(def_codes))
783 )
784 return clone