Coverage for src/usaspending/models/base_model.py: 100%

1# usaspendingapi/models/base.py 

2from typing import Optional, Dict, Any, List, TYPE_CHECKING 

3from weakref import ref 

4 

5if TYPE_CHECKING: 

6 from usaspending.client import USASpending 

7 

8 

9class BaseModel: 

10 """Base class for all models with fundamental behaviors.""" 

11 

12 def __init__(self, data: Dict[str, Any]): 

13 self._data = data or {} 

14 

15 @property 

16 def raw(self) -> Dict[str, Any]: 

17 """Get raw data dictionary.""" 

18 return self._data 

19 

20 def to_dict(self) -> Dict[str, Any]: 

21 """Convert to dictionary.""" 

22 return self._data 

23 

24 def get_value(self, keys: List[str], default: Any = None) -> Any: 

25 """Return the first non-None value from the given keys.""" 

26 if not isinstance(keys, list): 

27 keys = [keys] 

28 

29 if not isinstance(self._data, dict): 

30 raise TypeError("Empty object data") 

31 for key in keys: 

32 if key in self._data: 

33 value = self._data[key] 

34 if value is not None: # Check for non-None instead of truthiness 

35 return value 

36 return default 

37 

38 

39class ClientAwareModel(BaseModel): 

40 """Base class for all models that need API client access.""" 

41 

42 def __init__(self, data: Dict[str, Any], client: "USASpending"): 

43 super().__init__(data) 

44 self._client_ref = ref(client) # Weak reference prevents circular refs 

45 

46 @property 

47 def _client(self) -> Optional["USASpending"]: 

48 """Get client instance if still alive.""" 

49 return self._client_ref() if self._client_ref else None