Coverage for src/stable_yield_lab/core/repositories.py: 91%

1"""In-memory repositories for StableYieldLab data models.""" 

2 

3from __future__ import annotations 

4 

5from collections.abc import Iterable, Iterator 

6 

7import pandas as pd 

8 

9from .models import Pool, PoolReturn 

10 

11 

12class PoolRepository: 

13 """Lightweight in-memory collection with pandas export/import.""" 

14 

15 def __init__(self, pools: Iterable[Pool] | None = None) -> None: 

16 self._pools: list[Pool] = list(pools) if pools else [] 

17 

18 def add(self, pool: Pool) -> None: 

19 self._pools.append(pool) 

20 

21 def extend(self, items: Iterable[Pool]) -> None: 

22 self._pools.extend(items) 

23 

24 def filter( 

25 self, 

26 *, 

27 min_tvl: float = 0.0, 

28 min_base_apy: float = 0.0, 

29 chains: list[str] | None = None, 

30 auto_only: bool = False, 

31 stablecoins: list[str] | None = None, 

32 ) -> "PoolRepository": 

33 res: list[Pool] = [] 

34 for pool in self._pools: 

35 if pool.tvl_usd < min_tvl: 

36 continue 

37 if pool.base_apy < min_base_apy: 

38 continue 

39 if auto_only and not pool.is_auto: 

40 continue 

41 if chains and pool.chain not in chains: 41 ↛ 42line 41 didn't jump to line 42 because the condition on line 41 was never true

42 continue 

43 if stablecoins and pool.stablecoin not in stablecoins: 

44 continue 

45 res.append(pool) 

46 return PoolRepository(res) 

47 

48 def to_dataframe(self) -> pd.DataFrame: 

49 return pd.DataFrame([pool.to_dict() for pool in self._pools]) 

50 

51 def __len__(self) -> int: 

52 return len(self._pools) 

53 

54 def __iter__(self) -> Iterator[Pool]: 

55 return iter(self._pools) 

56 

57 

58class ReturnRepository: 

59 """Collection of :class:`PoolReturn` rows with pivot helper.""" 

60 

61 def __init__(self, rows: Iterable[PoolReturn] | None = None) -> None: 

62 self._rows: list[PoolReturn] = list(rows) if rows else [] 

63 

64 def extend(self, rows: Iterable[PoolReturn]) -> None: 

65 self._rows.extend(rows) 

66 

67 def to_timeseries(self) -> pd.DataFrame: 

68 if not self._rows: 68 ↛ 69line 68 didn't jump to line 69 because the condition on line 68 was never true

69 return pd.DataFrame() 

70 df = pd.DataFrame([row.to_dict() for row in self._rows]) 

71 df["timestamp"] = pd.to_datetime(df["timestamp"], utc=True) 

72 return df.pivot(index="timestamp", columns="name", values="period_return").sort_index() 

73 

74 

75__all__ = ["PoolRepository", "ReturnRepository"]