org-editor/doc_manager.py

import logging
import os
import sys
from datetime import datetime
from typing import List

import org_rw
from org_rw import OrgDoc, OrgTime

EXTENSIONS = (".org", ".org.txt")


def is_today(ot: OrgTime):
    now = datetime.now()
    return (
        (ot.time.year == now.year)
        and (ot.time.month == now.month)
        and (ot.time.day == now.day)
    )


class Agenda:
    def __init__(
        self,
        /,
        with_hour: List[org_rw.Headline],
        no_hour: List[org_rw.Headline],
    ):
        self.with_hour = with_hour
        self.no_hour = no_hour

    def print(self):
        for item in self.with_hour:
            print(item.scheduled.time, item.state, item.title)

        if len(self.with_hour) > 0:
            print("--------")

        for item in self.no_hour:
            print(item.scheduled.time, item.state, item.title)


class DocumentManager:
    docs: list[OrgDoc]

    def __init__(self, base_path: os.PathLike):
        self.base_path = base_path

    def load(self):
        top = os.path.abspath(self.base_path)

        docs = []

        for root, dirs, files in os.walk(top):
            # Prune dirs
            i = 0
            while i < len(dirs):
                if dirs[i].startswith(".git"):
                    del dirs[i]
                else:
                    i += 1

            # Process files
            for name in files:
                if all(map(lambda ext: not name.endswith(ext), EXTENSIONS)):
                    continue

                path = os.path.join(root, name)

                try:
                    doc = org_rw.load(open(path), extra_cautious=True)
                    docs.append(doc)
                except Exception as err:
                    import traceback

                    traceback.print_exc()
                    print(f"== On {path}")
                    sys.exit(1)

        logging.info("Loaded {} files".format(len(docs)))

        self.docs = docs

    def get_agenda(self) -> Agenda:
        headline_count = 0
        items_in_agenda = []
        now = datetime.now()

        for doc in self.docs:
            for hl in doc.getAllHeadlines():
                headline_count += 1

                if hl.scheduled and isinstance(hl.scheduled, OrgTime):
                    if is_today(hl.scheduled):
                        items_in_agenda.append(hl)
                    elif (hl.scheduled.time.to_datetime() < now) and hl.is_todo:
                        items_in_agenda.append(hl)

        logging.info("Read {} items".format(headline_count))
        logging.info("{} items in agenda today".format(len(items_in_agenda)))

        items_with_hour = [
            item
            for item in items_in_agenda
            if item.scheduled and is_today(item.scheduled) and item.scheduled.time.hour
        ]
        other_items = [
            item
            for item in items_in_agenda
            if not (
                item.scheduled and is_today(item.scheduled) and item.scheduled.time.hour
            )
        ]

        logging.info("{} items today for a specific hour".format(len(items_with_hour)))

        return Agenda(
            with_hour=sorted(items_with_hour, key=lambda x: x.scheduled.time),
            no_hour=other_items,
        )

    def get_notes(self, query) -> List[org_rw.Headline]:
        headline_count = 0
        t0 = datetime.now()
        notes = []
        query = [q.lower() for q in query]

        for doc in self.docs:
            for hl in doc.getAllHeadlines():
                headline_count += 1

                data = "\n".join(hl.get_contents("raw")).lower()
                if all([q in data for q in query]):
                    notes.append(hl)

        logging.info(
            "Filtered {} to {} items in {:.3f}s".format(
                headline_count, len(notes), (datetime.now() - t0).total_seconds()
            )
        )

        return notes