From ab10a5834ced32bc399af76154a2f85cb1a187a0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Lu=C3=ADs=20Murta?= <luis@murta.dev>
Date: Wed, 1 Dec 2021 18:41:07 +0000
Subject: [PATCH] Adds rudementary networth graph

---
 pfbudget/graph.py    | 44 ++++++++++++++++++++++++++++++++++++++++++++
 pfbudget/runnable.py |  4 +++-
 2 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/pfbudget/graph.py b/pfbudget/graph.py
index 3eb954d..a7ea9bc 100644
--- a/pfbudget/graph.py
+++ b/pfbudget/graph.py
@@ -177,3 +177,47 @@ def discrete(
         plt.savefig("graph.png")
     else:
         plt.show()
+
+
+def networth(
+    db: DBManager, args: dict, start: dt.date = dt.date.min, end: dt.date = dt.date.max
+):
+    transactions = db.get_daterange(start, end)
+    start, end = transactions[0].date, transactions[-1].date
+
+    accum = 0
+    monthly_networth = tuple(
+        (
+            month,
+            accum := sum(
+                transaction.value
+                for transaction in transactions
+                if transaction.original != "No"
+                and transaction.category not in pfbudget.categories.groups["investment"]
+                and month
+                <= transaction.date
+                <= month + dt.timedelta(days=monthrange(month.year, month.month)[1] - 1)
+            ) + accum
+        )
+        for month in [
+            month.date()
+            for month in rrule(
+                MONTHLY, dtstart=start.replace(day=1), until=end.replace(day=1)
+            )
+        ]
+    )
+
+    plt.figure(tight_layout=True)
+    plt.plot(
+        list(rrule(MONTHLY, dtstart=start.replace(day=1), until=end.replace(day=1))),
+        [
+            value for _, value in monthly_networth
+        ],
+        label="Total networth"
+    )
+    plt.grid()
+    plt.legend(loc="upper left")
+    if args["save"]:
+        plt.savefig("graph.png")
+    else:
+        plt.show()
diff --git a/pfbudget/runnable.py b/pfbudget/runnable.py
index 988be16..d1404d5 100644
--- a/pfbudget/runnable.py
+++ b/pfbudget/runnable.py
@@ -123,7 +123,7 @@ def argparser() -> argparse.ArgumentParser:
     p_graph.add_argument(
         "option",
         type=str,
-        choices=["monthly", "discrete"],
+        choices=["monthly", "discrete", "networth"],
         nargs="?",
         default="monthly",
         help="graph option help",
@@ -173,6 +173,8 @@ def graph(args):
         pfbudget.graph.monthly(DBManager(args.database), vars(args), start, end)
     elif args.option == "discrete":
         pfbudget.graph.discrete(DBManager(args.database), vars(args), start, end)
+    elif args.option == "networth":
+        pfbudget.graph.networth(DBManager(args.database), vars(args), start, end)
 
 
 def report(args):