Coverage for src/usaspending/resources/download_resource.py: 100%

1# src/usaspending/resources/download_resource.py 

2 

3from __future__ import annotations 

4from typing import TYPE_CHECKING, Optional 

5 

6from .base_resource import BaseResource 

7from ..logging_config import USASpendingLogger 

8# Import the manager and type aliases 

9from ..download.manager import DownloadManager, FileFormat 

10from ..download.job import DownloadJob 

11from ..models.download import DownloadStatus 

12 

13if TYPE_CHECKING: 

14 from ..client import USASpending 

15 

16logger = USASpendingLogger.get_logger(__name__) 

17 

18class DownloadResource(BaseResource): 

19 """Resource for award data download operations.""" 

20 

21 def __init__(self, client: USASpending): 

22 super().__init__(client) 

23 self._manager = DownloadManager(client) 

24 

25 def contract(self, award_id: str, file_format: FileFormat = "csv", destination_dir: Optional[str] = None) -> DownloadJob: 

26 """ 

27 Queue a download job for contract award data. 

28 

29 Args: 

30 award_id: The unique award identifier (e.g., CONT_AWD_...). 

31 file_format: Format of the file (csv, tsv, pstxt). 

32 destination_dir: Directory where the file will be saved (defaults to CWD). 

33 

34 Returns: 

35 A DownloadJob object. Use job.wait_for_completion() to block until finished. 

36 """ 

37 return self._manager.queue_download("contract", award_id, file_format, destination_dir) 

38 

39 def assistance(self, award_id: str, file_format: FileFormat = "csv", destination_dir: Optional[str] = None) -> DownloadJob: 

40 """ 

41 Queue a download job for assistance award data. 

42 

43 Args: 

44 award_id: The unique award identifier (e.g., ASST_NON_...). 

45 file_format: Format of the file (csv, tsv, pstxt). 

46 destination_dir: Directory where the file will be saved (defaults to CWD). 

47 

48 Returns: 

49 A DownloadJob object. Use job.wait_for_completion() to block until finished. 

50 """ 

51 return self._manager.queue_download("assistance", award_id, file_format, destination_dir) 

52 

53 def idv(self, award_id: str, file_format: FileFormat = "csv", destination_dir: Optional[str] = None) -> DownloadJob: 

54 """ 

55 Queue a download job for IDV (Indefinite Delivery Vehicle) award data. 

56 

57 Args: 

58 award_id: The unique award identifier (e.g., IDV_...). 

59 file_format: Format of the file (csv, tsv, pstxt). 

60 destination_dir: Directory where the file will be saved (defaults to CWD). 

61 

62 Returns: 

63 A DownloadJob object. Use job.wait_for_completion() to block until finished. 

64 """ 

65 return self._manager.queue_download("idv", award_id, file_format, destination_dir) 

66 

67 def status(self, file_name: str) -> DownloadStatus: 

68 """ 

69 Check the status of a specific download job directly via the API. 

70  

71 Note: Using DownloadJob.refresh_status() is generally preferred. 

72 

73 Args: 

74 file_name: The name of the file returned by the download request. 

75 

76 Returns: 

77 The DownloadStatus model representation. 

78 """ 

79 return self._manager.check_status(file_name)