Coverage for src/usaspending/models/download.py: 91%

1# src/usaspending/models/download.py 

2 

3from __future__ import annotations 

4from typing import Optional, Literal 

5from enum import Enum 

6 

7from ..utils.formatter import to_float, to_int 

8from .base_model import BaseModel 

9 

10AwardType = Literal["contract", "assistance", "idv"] 

11FileFormat = Literal["csv", "tsv", "pstxt"] 

12 

13class DownloadState(Enum): 

14 """Enumeration for download job states.""" 

15 PENDING = "pending" # Custom state before first API check 

16 READY = "ready" 

17 RUNNING = "running" 

18 FINISHED = "finished" 

19 FAILED = "failed" 

20 UNKNOWN = "unknown" # Custom state if API returns unexpected value 

21 

22class DownloadStatus(BaseModel): 

23 """Represents the status details of a download job returned by the API.""" 

24 

25 @property 

26 def file_name(self) -> Optional[str]: 

27 return self.get_value("file_name") 

28 

29 @property 

30 def message(self) -> Optional[str]: 

31 """A human readable error message if the status is failed, otherwise null.""" 

32 return self.get_value("message") 

33 

34 @property 

35 def seconds_elapsed(self) -> Optional[float]: 

36 """Time taken to generate the file or time taken so far.""" 

37 return to_float(self.get_value("seconds_elapsed")) 

38 

39 @property 

40 def api_status(self) -> DownloadState: 

41 """Current state of the request from the API.""" 

42 status_str = self.get_value("status") 

43 if status_str: 

44 try: 

45 return DownloadState(status_str) 

46 except ValueError: 

47 return DownloadState.UNKNOWN 

48 return DownloadState.UNKNOWN 

49 

50 @property 

51 def total_columns(self) -> Optional[int]: 

52 return to_int(self.get_value("total_columns")) 

53 

54 @property 

55 def total_rows(self) -> Optional[int]: 

56 return to_int(self.get_value("total_rows")) 

57 

58 @property 

59 def total_size_kb(self) -> Optional[float]: 

60 """Estimated file size in kilobytes.""" 

61 return to_float(self.get_value("total_size")) 

62 

63 @property 

64 def file_url(self) -> Optional[str]: 

65 """The URL for the file (relative path).""" 

66 return self.get_value("file_url") 

67 

68 def __repr__(self) -> str: 

69 return f"<DownloadStatus status='{self.api_status.value}' file='{self.file_name}'>"