From 2cbf00a9390b35105959d76e2f4dddf145c7255f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Lu=C3=ADs=20Murta?= <luis@murta.dev>
Date: Wed, 12 Apr 2023 18:20:57 +0100
Subject: [PATCH] CLI now allow --no-nulls flag on categorize

Null categorization, which is a major time slog, may now be disabled.
---
 pfbudget/__main__.py         |  6 ++++++
 pfbudget/cli/argparser.py    |  5 ++++-
 pfbudget/core/categorizer.py | 17 ++++++++++-------
 pfbudget/core/manager.py     |  2 +-
 4 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/pfbudget/__main__.py b/pfbudget/__main__.py
index 7a28ef5..a3fe169 100644
--- a/pfbudget/__main__.py
+++ b/pfbudget/__main__.py
@@ -25,6 +25,12 @@ if __name__ == "__main__":
             Interactive(Manager(db, verbosity)).start()
             exit()
 
+        case Operation.Categorize:
+            keys = {"no_nulls"}
+            assert args.keys() >= keys, f"missing {args.keys() - keys}"
+
+            params = [args["no_nulls"]]
+
         case Operation.Parse:
             keys = {"path", "bank", "creditcard"}
             assert args.keys() >= keys, f"missing {args.keys() - keys}"
diff --git a/pfbudget/cli/argparser.py b/pfbudget/cli/argparser.py
index 80f732d..e60cbfa 100644
--- a/pfbudget/cli/argparser.py
+++ b/pfbudget/cli/argparser.py
@@ -78,7 +78,10 @@ def argparser() -> argparse.ArgumentParser:
 
     # Automatic/manual categorization
     categorize = subparsers.add_parser("categorize").add_subparsers(required=True)
-    categorize.add_parser("auto").set_defaults(op=Operation.Categorize)
+    auto = categorize.add_parser("auto")
+    auto.set_defaults(op=Operation.Categorize)
+    auto.add_argument("--no-nulls", action="store_false")
+
     categorize.add_parser("manual").set_defaults(op=Operation.ManualCategorization)
 
     """
diff --git a/pfbudget/core/categorizer.py b/pfbudget/core/categorizer.py
index 4194379..b8b99df 100644
--- a/pfbudget/core/categorizer.py
+++ b/pfbudget/core/categorizer.py
@@ -16,6 +16,7 @@ class Categorizer:
         transactions: Sequence[t.BankTransaction],
         categories: Sequence[t.Category],
         tags: Sequence[t.Tag],
+        nullify: bool = True
     ):
         """Overarching categorization tool
 
@@ -28,14 +29,16 @@ class Categorizer:
             tags (Sequence[Tag]): currently available tags
         """
 
-        try:
-            null = next(cat for cat in categories if cat.name == "null")
-            print("Nullifying")
-            self._nullify(transactions, null)
+        if nullify:
+            try:
+                null = next(cat for cat in categories if cat.name == "null")
+                print("Nullifying")
+                self._nullify(transactions, null)
 
-            categories = [cat for cat in categories if cat.name != "null"]
-        except StopIteration:
-            print("Null category not defined")
+            except StopIteration:
+                print("Null category not defined")
+
+        categories = [cat for cat in categories if cat.name != "null"]
 
         self._rule_based_categories(transactions, categories)
         self._rule_based_tags(transactions, tags)
diff --git a/pfbudget/core/manager.py b/pfbudget/core/manager.py
index 9435e90..45b1982 100644
--- a/pfbudget/core/manager.py
+++ b/pfbudget/core/manager.py
@@ -97,7 +97,7 @@ class Manager:
                     )
                     categories = session.get(Category)
                     tags = session.get(Tag)
-                    Categorizer().rules(uncategorized, categories, tags)
+                    Categorizer().rules(uncategorized, categories, tags, params[0])
 
             case Operation.BankMod:
                 with self.db.session() as session: