Sort only in leaf categories

helper.py

@@ -6,18 +6,23 @@ import re
 from xdg import XDG_CONFIG_HOME, XDG_DATA_HOME
 from sqlalchemy import create_engine
 from sqlalchemy.orm import sessionmaker
-from models import Category, Transaction
+from models import Category
 
 try:
-    from yaml import CLoader as Loader, CDumper as Dumper
+    from yaml import CLoader as Loader
 except ImportError:
-    from yaml import Loader, Dumper
+    from yaml import Loader
 
 
 CONFIG_DIR = os.path.join(XDG_CONFIG_HOME, "schmeckels")
 DATA_DIR = os.path.join(XDG_DATA_HOME, "schmeckels")
 
 
+def get_list_of_bookable_categories(session):
+    categories = session.query(Category).all()
+    return [c for c in categories if c.is_child()]
+
+
 def create_dirs():
     for directory in [CONFIG_DIR, DATA_DIR]:
         try:

sort.py

@@ -1,26 +1,21 @@
 #! /usr/bin/env python3
-from sqlalchemy import create_engine
+from models import Transaction
-from sqlalchemy.orm import sessionmaker
+from helper import get_session, add_category, get_list_of_bookable_categories
-from sqlalchemy.ext.declarative import declarative_base
-from models import Category, Transaction
-from helper import get_session, add_category
 import sys
 import click
 
 from prompt_toolkit.completion import FuzzyWordCompleter
 from prompt_toolkit.shortcuts import prompt
 
-
 @click.command(name="sort")
 @click.option("--profile", "-p")
 def command(profile):
     session = get_session(profile)
-
+    categories = get_list_of_bookable_categories(session)
-    categories = session.query(Category).all()
     category_lookup = {c.full_name(): c.id for c in categories}
     category_names = FuzzyWordCompleter([c.full_name() for c in categories])
 
-    unsorted = session.query(Transaction).filter(Transaction.category_id == None).all()
+    unsorted = session.query(Transaction).filter(Transaction.category_id is None).all()
     print("Found {} unsorted transcations".format(len(unsorted)))
 
     for t in unsorted:
@@ -32,10 +27,16 @@ def command(profile):
         except KeyboardInterrupt:
             print("Goodbye")
             sys.exit(0)
+
         cat_id = category_lookup.get(select, None)
         if not cat_id:
             print(f"Creating new category '{select}'")
             cat_id = add_category(select, profile)
+            # Update the list of categories
+            categories = get_list_of_bookable_categories(session)
+            category_lookup = {c.full_name(): c.id for c in categories}
+            category_names = FuzzyWordCompleter([c.full_name() for c in categories])
+
         t.category_id = cat_id
         session.add(t)
         session.commit()