Coverage for src/usaspending/queries/recipients_search.py: 100%

1"""Recipients search query builder for USASpending recipient search endpoint.""" 

2 

3from __future__ import annotations 

4 

5from typing import Any, Optional, Literal, TYPE_CHECKING 

6 

7from usaspending.client import USASpending 

8from usaspending.queries.query_builder import QueryBuilder 

9from usaspending.logging_config import USASpendingLogger 

10 

11if TYPE_CHECKING: 

12 from usaspending.models.recipient import Recipient 

13 

14logger = USASpendingLogger.get_logger(__name__) 

15 

16# Valid award types for recipient searches 

17AwardType = Literal["all", "contracts", "grants", "loans", "direct_payments", "other_financial_assistance"] 

18# Valid sort fields 

19SortField = Literal["name", "duns", "amount"] 

20# Valid sort directions 

21SortDirection = Literal["asc", "desc"] 

22 

23 

24class RecipientsSearch(QueryBuilder["Recipient"]): 

25 """ 

26 Builds and executes recipient search queries, allowing for complex 

27 filtering on recipient data. This class follows a fluent interface pattern. 

28  

29 Supports filtering by keyword, award type, and sorting by various fields. 

30 """ 

31 

32 def __init__(self, client: USASpending): 

33 """ 

34 Initializes the RecipientsSearch query builder. 

35 

36 Args: 

37 client: The USASpending client instance. 

38 """ 

39 super().__init__(client) 

40 self._keyword: Optional[str] = None 

41 self._award_type: AwardType = "all" 

42 self._sort_field: SortField = "amount" 

43 self._sort_direction: SortDirection = "desc" 

44 

45 @property 

46 def _endpoint(self) -> str: 

47 """The API endpoint for this query.""" 

48 return "/recipient/" 

49 

50 def _clone(self) -> RecipientsSearch: 

51 """Creates an immutable copy of the query builder.""" 

52 clone = super()._clone() 

53 clone._keyword = self._keyword 

54 clone._award_type = self._award_type 

55 clone._sort_field = self._sort_field 

56 clone._sort_direction = self._sort_direction 

57 return clone 

58 

59 def _build_payload(self, page: int) -> dict[str, Any]: 

60 """Constructs the final API request payload.""" 

61 payload = { 

62 "page": page, 

63 "limit": self._get_effective_page_size(), 

64 "sort": self._sort_field, 

65 "order": self._sort_direction, 

66 "award_type": self._award_type, 

67 } 

68 

69 # Add keyword filter if provided 

70 if self._keyword: 

71 payload["keyword"] = self._keyword 

72 

73 return payload 

74 

75 def _transform_result(self, result: dict[str, Any]) -> "Recipient": 

76 """Transforms a single API result item into a Recipient model.""" 

77 from usaspending.models.recipient import Recipient 

78 if not result.get("recipient_id") and result.get("id"): 

79 result["recipient_id"] = result["id"] 

80 return Recipient(result, self._client) 

81 

82 def count(self) -> int: 

83 """ 

84 Get the total count of results using the dedicated count endpoint. 

85  

86 Uses the /v2/recipient/count/ endpoint which takes the same filters 

87 but returns just a count value. 

88 

89 Returns: 

90 The total number of matching recipients. 

91 """ 

92 logger.debug(f"{self.__class__.__name__}.count() called") 

93 

94 # Build payload for count endpoint (no pagination params needed) 

95 payload = { 

96 "award_type": self._award_type, 

97 } 

98 

99 # Add keyword filter if provided 

100 if self._keyword: 

101 payload["keyword"] = self._keyword 

102 

103 # Make API request to count endpoint 

104 count_endpoint = "/recipient/count/" 

105 response = self._client._make_request("POST", count_endpoint, json=payload) 

106 

107 total_count = response.get("count", 0) 

108 logger.info(f"{self.__class__.__name__}.count() = {total_count}") 

109 return total_count 

110 

111 # ========================================================================== 

112 # Filter Methods 

113 # ========================================================================== 

114 

115 def with_keyword(self, keyword: str) -> RecipientsSearch: 

116 """ 

117 Filter by recipient name, UEI, or DUNS keyword. 

118 

119 Args: 

120 keyword: The keyword to search for across recipient identifiers. 

121 

122 Returns: 

123 A new RecipientsSearch instance with the filter applied. 

124 """ 

125 clone = self._clone() 

126 clone._keyword = keyword.strip() if keyword else None 

127 return clone 

128 

129 def with_award_type(self, award_type: AwardType) -> RecipientsSearch: 

130 """ 

131 Filter by award type. 

132 

133 Args: 

134 award_type: The award type to filter by. 

135 

136 Returns: 

137 A new RecipientsSearch instance with the filter applied. 

138 """ 

139 clone = self._clone() 

140 clone._award_type = award_type 

141 return clone 

142 

143 def order_by(self, field: SortField, direction: SortDirection = "desc") -> RecipientsSearch: 

144 """ 

145 Set the sort field and direction. 

146 

147 Args: 

148 field: The field to sort by (name, duns, amount). 

149 direction: The sort direction (asc or desc). 

150 

151 Returns: 

152 A new RecipientsSearch instance with the sorting applied. 

153 """ 

154 clone = self._clone() 

155 clone._sort_field = field 

156 clone._sort_direction = direction 

157 return clone