1
0
fin-depo/fin_depo/data.py

58 lines
1.5 KiB
Python
Raw Permalink Normal View History

2024-06-02 15:24:53 +00:00
import abc
2024-05-29 20:29:15 +00:00
import dataclasses
2024-06-02 14:14:46 +00:00
import datetime
2024-06-02 15:24:53 +00:00
from collections.abc import Iterable, Mapping
2024-06-02 14:14:46 +00:00
from decimal import Decimal
import enforce_typing
2024-05-29 20:29:15 +00:00
from fin_defs import Asset
2024-06-02 14:14:46 +00:00
2024-05-29 20:29:15 +00:00
@enforce_typing.enforce_types
@dataclasses.dataclass
2024-06-02 15:24:53 +00:00
class Depo(abc.ABC):
2024-05-29 20:29:15 +00:00
name: str
updated_time: datetime.datetime
2024-06-02 15:24:53 +00:00
@abc.abstractmethod
def assets(self) -> Iterable[Asset]:
"""Returns the different assets managed by this depo."""
@abc.abstractmethod
def get_amount_of_asset(self, asset: Asset) -> Decimal:
"""Returns the amount of owned assets for all nested depos.
Must return 0 if depo does not contain the given asset.
"""
2024-06-02 14:14:46 +00:00
2024-05-29 20:29:15 +00:00
@enforce_typing.enforce_types
@dataclasses.dataclass
class DepoSingle(Depo):
2024-06-02 15:24:53 +00:00
_assets: Mapping[Asset, Decimal]
def assets(self) -> Iterable[Asset]:
return self._assets
def get_amount_of_asset(self, asset: Asset) -> Decimal:
return self._assets.get(asset, Decimal(0))
2024-05-29 20:29:15 +00:00
2024-06-02 14:14:46 +00:00
2024-05-29 20:29:15 +00:00
@enforce_typing.enforce_types
@dataclasses.dataclass
class DepoGroup(Depo):
nested: list[Depo]
2024-06-02 15:24:53 +00:00
def assets(self) -> Iterable[Asset]:
assets: list[Asset] = []
for nested_depo in self.nested:
assets.extend(nested_depo.assets())
return assets
def get_amount_of_asset(self, asset: Asset) -> Decimal:
summed: Decimal = Decimal(0)
for nested_depo in self.nested:
summed += nested_depo.get_amount_of_asset(asset)
del nested_depo
return summed