First attempt at balance!!

The output still needs colors, but it's functional. :D

2 files changed, 118 insertions, 0 deletions

diff --git a/ledger.py b/ledger.py
index 0e1e166..a418e65 100755
--- a/ledger.py
+++ b/ledger.py
@@ -3,6 +3,7 @@ from utils.args import get_arguments, test_args
 from utils.read_file import read_ledger
 from utils.register import print_register
 from utils.sort import sort_entries
+from utils.balance import print_balance
 
 def main():
     args = get_arguments()
@@ -23,6 +24,9 @@ def main():
     elif args.verb in ['register', 'reg', 'r']:
         print_register(result)
 
+    elif args.verb in ['balance', 'bal', 'b']:
+        print_balance(result)
+
 
 if __name__ == '__main__':
     main()
diff --git a/utils/balance.py b/utils/balance.py
new file mode 100644
index 0000000..64fa3e3
--- /dev/null
+++ b/utils/balance.py
@@ -0,0 +1,114 @@
+from typing import List
+
+from utils.read_file import entry
+from utils.register import currencies, complete_prices
+
+
+class tree:
+    def __init__(self, name: str, value: currencies = None) -> None:
+        self.name = name
+        self.value = value
+        self.children = []
+
+    
+    def compute_value(self):
+        if self.value is not None:
+            return
+
+        self.value = currencies()
+        for child in reversed(self.children):
+            child.compute_value()
+            for price in child.value:
+                self.value.add_money(price)
+
+
+    def add_child(self, name: str, value: currencies = None):
+        self.children.append(tree(name, value))
+
+
+    def add_value(self, price: str):
+        if self.value is None:
+            self.value = currencies()
+
+        self.value.add_money(price)
+
+
+    def get_child_with_name(self, name: str):
+        for child in self.children:
+            if name == child.name:
+                return child
+
+        return None
+
+
+    def __contains__(self, name: str) -> bool:
+        for child in self.children:
+            if name == child.name:
+                return True
+            
+        return False
+
+    
+    def __str__(self, level=0, ignore_level = False) -> str:
+        result = ''
+        if not ignore_level:
+            for price in self.value:
+                result += f'{price:>20}\n'
+
+            result = result.rstrip()
+            result += '  ' + ' '*(2 * level)
+
+        if len(self.children) == 1:
+            result += self.name + ':'
+            for child in self.children:
+                result += child.__str__(level + 1, ignore_level = True)
+
+        else:
+            result += self.name + '\n'
+            for child in self.children:
+                result += child.__str__(level + 1)
+
+        return result
+
+
+class accounts_tree:
+    def __init__(self) -> None:
+        self.my_accounts = tree(name='root')
+
+    def add_transaction(self, name_acc: str, price: str):
+        search = self.my_accounts
+
+        for name in name_acc.split(':'):
+            if name not in search:
+                search.add_child(name)
+
+            search = search.get_child_with_name(name)
+
+        search.add_value(price)
+
+
+    def __str__(self) -> str:
+        self.my_accounts.compute_value()
+
+        result = ''
+        for child in self.my_accounts.children:
+            result += str(child)
+
+        result += '--------------------\n'
+        for price in self.my_accounts.value:
+            result += f'{price:>20}\n'
+
+        return result.rstrip()
+
+
+
+def print_balance(my_entries: List[entry]):
+    my_tree = accounts_tree()
+
+    for ent in my_entries:
+        complete_prices(ent)
+
+        for trans in ent.transactions:
+            my_tree.add_transaction(trans[0], str(trans[1]))
+
+    print(my_tree)