Coverage for src/usaspending/models/location.py: 97%

1from __future__ import annotations 

2from typing import Dict, Any, Optional 

3from titlecase import titlecase 

4from ..utils.formatter import contracts_titlecase 

5from .base_model import BaseModel 

6 

7 

8class Location(BaseModel): 

9 """Location model for USASpending data.""" 

10 

11 def __init__(self, data: Dict[str, Any], client=None): 

12 """Initialize Location. Client parameter is ignored for compatibility.""" 

13 super().__init__(data) 

14 

15 # simple direct fields -------------------------------------------------- 

16 @property 

17 def address_line1(self) -> Optional[str]: 

18 return self._format_location_string_property(self.get_value(["address_line1"])) 

19 

20 @property 

21 def address_line2(self) -> Optional[str]: 

22 return self._format_location_string_property(self.get_value(["address_line2"])) 

23 

24 @property 

25 def address_line3(self) -> Optional[str]: 

26 return self._format_location_string_property(self.get_value(["address_line3"])) 

27 

28 @property 

29 def city_name(self) -> Optional[str]: 

30 city_name = self.get_value(["city_name", "city"]) 

31 if not isinstance(city_name, str): 

32 return None 

33 return titlecase(city_name) 

34 

35 @property 

36 def city(self) -> Optional[str]: 

37 return self.city_name 

38 

39 @property 

40 def state_name(self) -> Optional[str]: 

41 state_name = titlecase(self.get_value(["state_name", "state"])) 

42 if not isinstance(state_name, str): 

43 return None 

44 return titlecase(state_name) 

45 

46 @property 

47 def country_name(self) -> Optional[str]: 

48 country = self._format_location_string_property( 

49 self.get_value(["country_name"]) 

50 ) 

51 if country and country.lower() == "usa": 

52 country = "USA" 

53 return country 

54 

55 @property 

56 def zip4(self) -> Optional[str]: 

57 return self.get_value(["zip4"]) 

58 

59 @property 

60 def county_name(self) -> Optional[str]: 

61 county_name = titlecase(self.get_value(["county_name", "county"])) 

62 if not isinstance(county_name, str): 

63 return None 

64 return county_name 

65 

66 @property 

67 def county_code(self) -> Optional[str]: 

68 return self.get_value(["county_code"]) 

69 

70 @property 

71 def congressional_code(self) -> Optional[str]: 

72 return self.get_value(["congressional_code","district"]) 

73 

74 @property 

75 def foreign_province(self) -> Optional[str]: 

76 return self.get_value(["foreign_province"]) 

77 

78 @property 

79 def foreign_postal_code(self) -> Optional[str]: 

80 return self.get_value(["foreign_postal_code"]) 

81 

82 # dual-source fields ---------------------------------------------------- 

83 @property 

84 def state_code(self) -> Optional[str]: 

85 return self.get_value(["state_code", "Place of Performance State Code"]) 

86 

87 @property 

88 def country_code(self) -> Optional[str]: 

89 return self.get_value( 

90 ["location_country_code", "Place of Performance Country Code"] 

91 ) 

92 

93 @property 

94 def zip5(self) -> Optional[str]: 

95 val = self.get_value(["zip5", "Place of Performance Zip5"]) 

96 return str(val) if val is not None else "" 

97 

98 # convenience ----------------------------------------------------------- 

99 @property 

100 def district(self) -> Optional[str]: 

101 pieces = [p for p in (self.state_code, self.congressional_code) if p] 

102 return "-".join(pieces) or "" 

103 

104 @property 

105 def formatted_address(self) -> Optional[str]: 

106 lines: list[str] = [ 

107 line 

108 for line in (self.address_line1, self.address_line2, self.address_line3) 

109 if line 

110 ] 

111 trailing = [p for p in (self.city, self.state_code, self.zip5) if p] 

112 if trailing: 

113 lines.append(", ".join(trailing)) 

114 if self.country_name: 

115 lines.append(self.country_name) 

116 return "\n".join(lines) or None 

117 

118 def _format_location_string_property(self, text: str) -> Optional[str]: 

119 """Format a location string with string check.""" 

120 if not isinstance(text, str): 

121 return None 

122 return contracts_titlecase(text.strip()) 

123 

124 def __repr__(self) -> str: 

125 return f"<Location {self.city or '?'} {self.state_code or ''} {self.country_code or ''}>"