Tools/time_recoder/time_recoder_gui/time_recorder_task_handler.py

import sqlalchemy
from sqlalchemy import select
from sqlalchemy.orm import sessionmaker
from time_recoder.time_recorder_database.create_db import ENGINE, create_root_task
from time_recoder.time_recorder_database.db_models import Task

Session = sessionmaker(bind=ENGINE)


class TaskHandler:

    def __init__(self):
        self.session = Session()
        self.root_task = self.get_root_task()

    def get_root_task(self):
        task: Task
        try:
            task = self.session.query(Task).filter(Task.parent_task_id == -1).one()
        except:
            create_root_task()
            task = self.session.query(Task).filter(Task.parent_task_id == -1).one()
        return task

    def add_task(self, parent_task_id: int, name: str):
        try:
            new_task = Task(parent_task_id=parent_task_id, name=name)
            self.session.add(new_task)
            self.session.commit()
        except:
            print("Error: Couln't commit new Task to DB. Please close any other DB connections")

    """ fancy recursion tree ♥ """
    def get_task_tree(self, parent_task: Task):
        child_tasks = self.session.query(Task).filter(Task.parent_task_id == parent_task.id).all()
        return {'task': parent_task,
                'child_tasks': [self.get_task_tree(task) for task in child_tasks]}

    def close_session(self):
        self.session.close()

    def task_tree_as_string_list(self, task_tree, level_offset=0):
        string_level_prefix = ''
        for _ in range(0, level_offset):
            string_level_prefix += '  '
        intended_name = string_level_prefix + task_tree['task'].name
        level_offset += 2
        child_tree_list = [self.task_tree_as_string_list(task_tree, level_offset) for task_tree in
                           task_tree['child_tasks']]
        flatten_child_tree_list = [item for sublist in child_tree_list for item in sublist]
        return [intended_name] + flatten_child_tree_list

    def task_tree_as_id_list(self, task_tree):
        child_tree_list = [self.task_tree_as_id_list(task_tree) for task_tree in task_tree['child_tasks']]
        root_id = task_tree['task'].id
        flatten_child_tree_list = [item for sublist in child_tree_list for item in sublist]
        return [root_id] + flatten_child_tree_list

    def get_task_tree_name_strings(self, task):
        return self.task_tree_as_string_list(self.get_task_tree(task))

    def get_task_tree_ids(self, task):
        return self.task_tree_as_id_list(self.get_task_tree(task))

    def get_task_id(self, task_name: str):
        return self.get_task_tree_ids(self.root_task)[self.get_task_tree_name_strings(self.root_task).index(task_name)]

    def get_task_by_id(self, task_id: int):
        return self.session.query(Task).filter(Task.id == task_id).one()

    def get_task_branch_as_string(self, task):
        parent_task_id = task.parent_task_id
        if parent_task_id == -1:
            return task.name
        else:
            return self.get_task_branch_as_string(self.get_task_by_id(parent_task_id)) + '/' + task.name