diff --git a/helper.py b/helper.py index 28c3f0f..0308855 100644 --- a/helper.py +++ b/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: diff --git a/sort.py b/sort.py index 0640f25..b17a675 100644 --- a/sort.py +++ b/sort.py @@ -1,26 +1,21 @@ #! /usr/bin/env python3 -from sqlalchemy import create_engine -from sqlalchemy.orm import sessionmaker -from sqlalchemy.ext.declarative import declarative_base -from models import Category, Transaction -from helper import get_session, add_category +from models import Transaction +from helper import get_session, add_category, get_list_of_bookable_categories 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 = session.query(Category).all() + 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]) - 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()