usaspending-orm 0.6.0__py3-none-any.whl

usaspending/queries/recipient_query.py

@@ -0,0 +1,66 @@
+from typing import TYPE_CHECKING
+from .single_resource_base import SingleResourceBase
+from ..exceptions import ValidationError
+from ..client import USASpendingClient
+from ..logging_config import USASpendingLogger
+
+if TYPE_CHECKING:
+    from ..models.recipient import Recipient
+
+logger = USASpendingLogger.get_logger(__name__)
+
+
+class RecipientQuery(SingleResourceBase):
+    """Retrieve a single-recipient"""
+
+    def __init__(self, client: USASpendingClient):
+        super().__init__(client)
+
+    @property
+    def _endpoint(self) -> str:
+        """Base endpoint for single recipient retrieval."""
+        return "/recipient/"
+
+    def find_by_id(self, recipient_id: str) -> "Recipient":
+        """Filter by unique recipient identifier."""
+        if not recipient_id:
+            raise ValidationError("recipient_id is required")
+
+        # Make API request
+        response = self._get_resource(recipient_id)
+
+        # Create model instance
+        from ..models.recipient import Recipient
+
+        return Recipient(response, client=self._client)
+
+    def _clean_recipient_id(self, recipient_id: str) -> str:
+        """Clean recipient ID format.
+
+        Handles cases like "abc123-['C','R']" -> "abc123-R"
+        """
+        import re
+
+        # Pattern for list suffix: -[...]
+        pattern = r"^(.+?)-\[\s*([^\]]+)\]$"
+        match = re.match(pattern, recipient_id)
+
+        if not match:
+            return recipient_id
+
+        base = match.group(1)
+        tokens = match.group(2)
+
+        # Extract and clean tokens
+        clean_tokens = []
+        for token in tokens.split(","):
+            clean_token = token.strip().strip("'\"").upper()
+            if clean_token:
+                clean_tokens.append(clean_token)
+
+        # Prefer 'R' if present, otherwise first token
+        suffix = (
+            "R" if "R" in clean_tokens else (clean_tokens[0] if clean_tokens else "")
+        )
+
+        return f"{base}-{suffix}" if suffix else base

usaspending/queries/recipients_search.py

@@ -0,0 +1,167 @@
+"""Recipients search query builder for USASpending recipient search endpoint."""
+
+from __future__ import annotations
+
+from typing import Any, Optional, Literal, TYPE_CHECKING
+
+from ..client import USASpendingClient
+from usaspending.queries.query_builder import QueryBuilder
+from usaspending.logging_config import USASpendingLogger
+
+if TYPE_CHECKING:
+    from usaspending.models.recipient import Recipient
+
+logger = USASpendingLogger.get_logger(__name__)
+
+# Valid award types for recipient searches
+AwardType = Literal[
+    "all",
+    "contracts",
+    "grants",
+    "loans",
+    "direct_payments",
+    "other_financial_assistance",
+]
+# Valid sort fields
+SortField = Literal["name", "duns", "amount"]
+# Valid sort directions
+SortDirection = Literal["asc", "desc"]
+
+
+class RecipientsSearch(QueryBuilder["Recipient"]):
+    """
+    Builds and executes recipient search queries, allowing for complex
+    filtering on recipient data. This class follows a fluent interface pattern.
+
+    Supports filtering by keyword, award type, and sorting by various fields.
+    """
+
+    def __init__(self, client: USASpendingClient):
+        """
+        Initializes the RecipientsSearch query builder.
+
+        Args:
+            client: The USASpending client instance.
+        """
+        super().__init__(client)
+        self._keyword: Optional[str] = None
+        self._award_type: AwardType = "all"
+        self._sort_field: SortField = "amount"
+        self._sort_direction: SortDirection = "desc"
+
+    @property
+    def _endpoint(self) -> str:
+        """The API endpoint for this query."""
+        return "/recipient/"
+
+    def _clone(self) -> RecipientsSearch:
+        """Creates an immutable copy of the query builder."""
+        clone = super()._clone()
+        clone._keyword = self._keyword
+        clone._award_type = self._award_type
+        clone._sort_field = self._sort_field
+        clone._sort_direction = self._sort_direction
+        return clone
+
+    def _build_payload(self, page: int) -> dict[str, Any]:
+        """Constructs the final API request payload."""
+        payload = {
+            "page": page,
+            "limit": self._get_effective_page_size(),
+            "sort": self._sort_field,
+            "order": self._sort_direction,
+            "award_type": self._award_type,
+        }
+
+        # Add keyword filter if provided
+        if self._keyword:
+            payload["keyword"] = self._keyword
+
+        return payload
+
+    def _transform_result(self, result: dict[str, Any]) -> "Recipient":
+        """Transforms a single API result item into a Recipient model."""
+        from usaspending.models.recipient import Recipient
+
+        if not result.get("recipient_id") and result.get("id"):
+            result["recipient_id"] = result["id"]
+        return Recipient(result, self._client)
+
+    def count(self) -> int:
+        """
+        Get the total count of results using the dedicated count endpoint.
+
+        Uses the /v2/recipient/count/ endpoint which takes the same filters
+        but returns just a count value.
+
+        Returns:
+            The total number of matching recipients.
+        """
+        logger.debug(f"{self.__class__.__name__}.count() called")
+
+        # Build payload for count endpoint (no pagination params needed)
+        payload = {
+            "award_type": self._award_type,
+        }
+
+        # Add keyword filter if provided
+        if self._keyword:
+            payload["keyword"] = self._keyword
+
+        # Make API request to count endpoint
+        count_endpoint = "/recipient/count/"
+        response = self._client._make_request("POST", count_endpoint, json=payload)
+
+        total_count = response.get("count", 0)
+        logger.info(f"{self.__class__.__name__}.count() = {total_count}")
+        return total_count
+
+    # ==========================================================================
+    # Filter Methods
+    # ==========================================================================
+
+    def keyword(self, keyword: str) -> RecipientsSearch:
+        """
+        Filter by recipient name, UEI, or DUNS keyword.
+
+        Args:
+            keyword: The keyword to search for across recipient identifiers.
+
+        Returns:
+            A new RecipientsSearch instance with the filter applied.
+        """
+        clone = self._clone()
+        clone._keyword = keyword.strip() if keyword else None
+        return clone
+
+    def award_type(self, award_type: AwardType) -> RecipientsSearch:
+        """
+        Filter by award type.
+
+        Args:
+            award_type: The award type to filter by.
+
+        Returns:
+            A new RecipientsSearch instance with the filter applied.
+        """
+        clone = self._clone()
+        clone._award_type = award_type
+        return clone
+
+    def order_by(
+        self, field: SortField, direction: SortDirection = "desc"
+    ) -> RecipientsSearch:
+        """
+        Set the sort field and direction.
+
+        Args:
+            field: The field to sort by (name, duns, amount).
+            direction: The sort direction (asc or desc).
+
+        Returns:
+            A new RecipientsSearch instance with the sorting applied.
+        """
+        clone = self._clone()
+        clone._sort_field = field
+        clone._sort_direction = direction
+        return clone

usaspending/queries/single_resource_base.py

@@ -0,0 +1,83 @@
+from ..exceptions import ValidationError
+from ..client import USASpendingClient
+from typing import Any
+from abc import ABC, abstractmethod
+from ..logging_config import USASpendingLogger
+
+logger = USASpendingLogger.get_logger(__name__)
+
+
+class SingleResourceBase(ABC):
+    """
+    Base class for retrieving single resources like awards or recipients,
+    anything with a simple GET request, generally.
+    """
+
+    def __init__(self, client: USASpendingClient):
+        self._client = client
+
+    @property
+    @abstractmethod
+    def _endpoint(self) -> str:
+        """Base endpoint for single resource retrieval."""
+        pass
+
+    @abstractmethod
+    def find_by_id(self, resource_id: str) -> Any:
+        """Filter by unique resource identifier."""
+        pass
+
+    def _get_resource(self, resource_id: str) -> dict:
+        """Retrieve a single resource by ID."""
+        if (
+            not resource_id
+            or not isinstance(resource_id, str)
+            or not resource_id.strip()
+        ):
+            raise ValidationError("A non-empty resource_id string is required")
+
+        # Clean recipient ID
+        cleaned_resource_id = self._clean_resource_id(resource_id)
+
+        if not cleaned_resource_id:
+            raise ValidationError(
+                "No resource id found after cleaning. Original: %s. Cleaned: %s",
+                resource_id,
+                cleaned_resource_id,
+            )
+
+        # Construct valid endpoint
+        endpoint = self._construct_endpoint(cleaned_resource_id)
+
+        # Make API request
+        response = self._client._make_request("GET", endpoint)
+
+        # Validate response
+        if response is None:
+            raise ValidationError(
+                f"API request for {endpoint} returned None. This may be due to caching issues."
+            )
+
+        if not isinstance(response, dict):
+            raise ValidationError(
+                f"API request for {endpoint} returned invalid response type: {type(response)}. Expected dict."
+            )
+
+        return response
+
+    def _construct_endpoint(self, resource_id: str) -> str:
+        """Construct the full endpoint URL for a specific resource ID."""
+        return f"{self._endpoint}{resource_id}/"
+
+    def _clean_resource_id(self, resource_id: str) -> str:
+        """Very basic resource ID cleaning.
+
+        More complex logic is implemented in specific resource classes.
+
+        Args:
+            resource_id: The raw resource ID string
+
+        Returns:
+            Cleaned resource ID string
+        """
+        return str(resource_id).strip()

usaspending/queries/spending_search.py

@@ -0,0 +1,261 @@
+"""Spending search query builder for USASpending spending by category endpoints."""
+
+from __future__ import annotations
+
+from typing import Any, Optional, Literal
+
+from ..client import USASpendingClient
+from usaspending.exceptions import ValidationError
+from usaspending.models.spending import Spending
+from usaspending.models.recipient_spending import RecipientSpending
+from usaspending.models.district_spending import DistrictSpending
+from usaspending.models.state_spending import StateSpending
+from usaspending.queries.query_builder import SearchQueryBuilder
+from usaspending.logging_config import USASpendingLogger
+from usaspending.queries.filters import (
+    SimpleListFilter,
+)
+
+# Import award type codes from models
+# These are defined by USASpending.gov and represent different categories of awards
+
+
+logger = USASpendingLogger.get_logger(__name__)
+
+SpendingLevel = Literal["transactions", "awards", "subawards"]
+SpendingCategory = Literal["recipient", "district", "state"]
+
+
+class SpendingSearch(SearchQueryBuilder["Spending"]):
+    """
+    Builds and executes spending by category search queries, allowing for complex
+    filtering on spending data. This class follows a fluent interface pattern.
+
+    Supports both recipient and district spending searches with configurable
+    spending levels (transactions, awards, subawards).
+    """
+
+    def __init__(self, client: USASpendingClient):
+        """
+        Initializes the SpendingSearch query builder.
+
+        Args:
+            client: The USASpending client instance.
+        """
+        super().__init__(client)
+        self._category: Optional[SpendingCategory] = None
+        self._spending_level: SpendingLevel = "transactions"
+        self._subawards: bool = False
+
+    @property
+    def _endpoint(self) -> str:
+        """The API endpoint for this query."""
+        if self._category == "recipient":
+            return "/search/spending_by_category/recipient/"
+        elif self._category == "district":
+            return "/search/spending_by_category/district/"
+        elif self._category == "state":
+            return "/search/spending_by_category/state_territory/"
+        else:
+            raise ValidationError(
+                "Category must be set. Use .by_recipient(), .by_district(), or .by_state() method."
+            )
+
+    def _clone(self) -> SpendingSearch:
+        """Creates an immutable copy of the query builder."""
+        clone = super()._clone()
+        clone._category = self._category
+        clone._spending_level = self._spending_level
+        clone._subawards = self._subawards
+        return clone
+
+    def _build_payload(self, page: int) -> dict[str, Any]:
+        """Constructs the final API request payload from the filter objects."""
+
+        if self._category is None:
+            raise ValidationError(
+                "Category must be set. Use .by_recipient(), .by_district(), or .by_state() method."
+            )
+
+        final_filters = self._aggregate_filters()
+
+        payload = {
+            "filters": final_filters,
+            "category": self._category,
+            "limit": self._get_effective_page_size(),
+            "page": page,
+            "spending_level": self._spending_level,
+        }
+
+        # Add deprecated subawards field if needed
+        if self._subawards:
+            payload["subawards"] = self._subawards
+
+        return payload
+
+    def _transform_result(self, result: dict[str, Any]) -> Spending:
+        """Transforms a single API result item into appropriate Spending model."""
+        # Add category info to result data for model initialization
+        result_with_category = {
+            **result,
+            "category": self._category,
+            "spending_level": self._spending_level,
+        }
+
+        if self._category == "recipient":
+            return RecipientSpending(result_with_category, self._client)
+        elif self._category == "district":
+            return DistrictSpending(result_with_category, self._client)
+        elif self._category == "state":
+            return StateSpending(result_with_category, self._client)
+        else:
+            return Spending(result_with_category, self._client)
+
+    def count(self) -> int:
+        """
+        Get the total count of results by iterating through pages.
+
+        Respects pagination constraints like limit() and max_pages() to match
+        the behavior of iteration. The spending by category endpoints don't
+        have a total count in page_metadata, so we fetch pages and count results.
+
+        Returns:
+            The total number of matching spending records, up to any set limits.
+        """
+        logger.debug(f"{self.__class__.__name__}.count() called")
+
+        # Early return for zero or negative limits
+        if self._total_limit is not None and self._total_limit <= 0:
+            logger.info(
+                f"{self.__class__.__name__}.count() = 0 (limit: {self._total_limit})"
+            )
+            return 0
+
+        total_count = 0
+        page = 1
+        pages_fetched = 0
+
+        while True:
+            # Check if we've reached the max pages limit
+            if self._max_pages and pages_fetched >= self._max_pages:
+                logger.debug(f"Max pages limit ({self._max_pages}) reached")
+                break
+
+            response = self._execute_query(page)
+            results = response.get("results", [])
+
+            # Count items, but respect total_limit
+            items_to_count = len(results)
+            if self._total_limit is not None:
+                remaining = self._total_limit - total_count
+                items_to_count = min(items_to_count, remaining)
+
+            total_count += items_to_count
+
+            # Stop if we've reached our limit
+            if self._total_limit is not None and total_count >= self._total_limit:
+                logger.debug(f"Total limit of {self._total_limit} items reached")
+                break
+
+            # Check if there are more pages
+            page_metadata = response.get("page_metadata", {})
+            has_next = page_metadata.get("hasNext", False)
+
+            if not has_next or not results:
+                break
+
+            page += 1
+            pages_fetched += 1
+
+        logger.info(f"{self.__class__.__name__}.count() = {total_count}")
+        return total_count
+
+    # ==========================================================================
+    # Category Selection Methods
+    # ==========================================================================
+
+    def by_recipient(self) -> SpendingSearch:
+        """
+        Configure search to return spending grouped by recipient.
+
+        Returns:
+            A new SpendingSearch instance configured for recipient spending.
+        """
+        clone = self._clone()
+        clone._category = "recipient"
+        return clone
+
+    def by_district(self) -> SpendingSearch:
+        """
+        Configure search to return spending grouped by congressional district.
+
+        Returns:
+            A new SpendingSearch instance configured for district spending.
+        """
+        clone = self._clone()
+        clone._category = "district"
+        return clone
+
+    def by_state(self) -> SpendingSearch:
+        """
+        Configure search to return spending grouped by state/territory.
+
+        Returns:
+            A new SpendingSearch instance configured for state spending.
+        """
+        clone = self._clone()
+        clone._category = "state"
+        return clone
+
+    # ==========================================================================
+    # Spending Level Configuration
+    # ==========================================================================
+
+    def spending_level(self, level: SpendingLevel) -> SpendingSearch:
+        """
+        Set the spending level for data aggregation.
+
+        Args:
+            level: The spending level - "transactions", "awards", or "subawards"
+
+        Returns:
+            A new SpendingSearch instance with the spending level configured.
+        """
+        clone = self._clone()
+        clone._spending_level = level
+        return clone
+
+    def subawards_only(self, enabled: bool = True) -> SpendingSearch:
+        """
+        Enable subawards search (deprecated parameter).
+
+        Args:
+            enabled: Whether to search subawards instead of prime awards
+
+        Returns:
+            A new SpendingSearch instance with subawards flag set.
+        """
+        clone = self._clone()
+        clone._subawards = enabled
+        return clone
+
+    # ==========================================================================
+    # Filter Methods (same as AwardsSearch)
+    # ==========================================================================
+
+    def recipient_id(self, recipient_id: str) -> SpendingSearch:
+        """
+        Filter by specific recipient ID.
+
+        Args:
+            recipient_id: Unique identifier for the recipient.
+
+        Returns:
+            A new SpendingSearch instance with the filter applied.
+        """
+        clone = self._clone()
+        clone._filter_objects.append(
+            SimpleListFilter(key="recipient_id", values=[recipient_id])
+        )
+        return clone
+