From f74945d6bf39fc206df2e8124ea42947166ad187 Mon Sep 17 00:00:00 2001
From: Jon Michael Aanes <jonjmaa@gmail.com>
Date: Sun, 2 Jun 2024 17:24:53 +0200
Subject: [PATCH] Depo methods

---
 fin_depo/_version.py    |  2 +-
 fin_depo/data.py        | 37 ++++++++++++++++++++++++++++++++++---
 fin_depo/defi_kucoin.py |  3 ++-
 setup.py                |  3 +++
 4 files changed, 40 insertions(+), 5 deletions(-)

diff --git a/fin_depo/_version.py b/fin_depo/_version.py
index 34da6b8..10939f0 100644
--- a/fin_depo/_version.py
+++ b/fin_depo/_version.py
@@ -1 +1 @@
-__version__ = '0.1.2'
\ No newline at end of file
+__version__ = '0.1.2'
diff --git a/fin_depo/data.py b/fin_depo/data.py
index 002412b..ceef6c4 100644
--- a/fin_depo/data.py
+++ b/fin_depo/data.py
@@ -1,6 +1,7 @@
+import abc
 import dataclasses
 import datetime
-from collections.abc import Mapping
+from collections.abc import Iterable, Mapping
 from decimal import Decimal
 
 import enforce_typing
@@ -9,18 +10,48 @@ from fin_defs import Asset
 
 @enforce_typing.enforce_types
 @dataclasses.dataclass
-class Depo:
+class Depo(abc.ABC):
     name: str
     updated_time: datetime.datetime
 
+    @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.
+        """
+
 
 @enforce_typing.enforce_types
 @dataclasses.dataclass
 class DepoSingle(Depo):
-    assets: Mapping[Asset, Decimal]
+    _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))
 
 
 @enforce_typing.enforce_types
 @dataclasses.dataclass
 class DepoGroup(Depo):
     nested: list[Depo]
+
+    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
diff --git a/fin_depo/defi_kucoin.py b/fin_depo/defi_kucoin.py
index f6f36cd..55d895e 100644
--- a/fin_depo/defi_kucoin.py
+++ b/fin_depo/defi_kucoin.py
@@ -25,7 +25,8 @@ class KucoinDepoFetcher:
             asset = fin_defs.WELL_KNOWN_SYMBOLS[account_data['currency']]
             balance = Decimal(account_data['balance'])
             assets_for_account_type = assets_by_account_type.get(
-                account_data['type'], {},
+                account_data['type'],
+                {},
             )
             assets_for_account_type[asset] = (
                 assets_for_account_type.get(asset, Decimal(0)) + balance
diff --git a/setup.py b/setup.py
index e77c2db..ef3b94c 100644
--- a/setup.py
+++ b/setup.py
@@ -15,6 +15,7 @@ PACKAGE_NAME = 'fin_depo'
 with open('README.md') as f:
     readme = f.read()
 
+
 def parse_version_file(text: str) -> str:
     match = re.match(r'^__version__\s*=\s*(["\'])([\d\.]+)\1$', text)
     if match is None:
@@ -22,9 +23,11 @@ def parse_version_file(text: str) -> str:
         raise Exception(msg)
     return match.group(2)
 
+
 with open(PACKAGE_NAME + '/_version.py') as f:
     version = parse_version_file(f.read())
 
+
 def parse_requirements(text: str) -> list[str]:
     return text.strip().split('\n')