From: Dustin Walde <redacted>
Date: Mon, 13 Nov 2023 01:27:18 +0000 (-0800)
Subject: Add auto_document plugin
X-Git-Url: https://git.walde.dev/?a=commitdiff_plain;h=83e3128cc566454aaf86c26ad7229356e332e66c;p=beanbeanbean

Add auto_document plugin
---

diff --git a/src/beanbeanbean/auto_document.py b/src/beanbeanbean/auto_document.py
new file mode 100644
index 0000000..80c36c5
--- /dev/null
+++ b/src/beanbeanbean/auto_document.py
@@ -0,0 +1,103 @@
+from dataclasses import dataclass
+from datetime import date
+
+from beancount.core.data import (
+    Custom,
+    Entries,
+    Balance,
+    Transaction,
+)
+from beancount.loader import LoadError
+
+from .utils import make_error
+
+from typing import List, Optional, Union
+
+__plugins__ = ('auto_document',)
+
+
+@dataclass
+class AutoDoc:
+    start_date: date
+    end_date: Optional[date]
+    account: str
+    document: str
+
+
+def custom_to_autodoc(entry: Custom) -> Union[AutoDoc,LoadError]:
+    start_date = entry.date
+    end_date = None
+    account = None
+    document = None
+    for item in entry.values:
+        if item.dtype is str and document is None:
+            document = item.value
+        elif item.dtype == '<AccountDummy>':
+            account = item.value
+        elif item.dtype is date:
+            end_date = item.value
+
+    if account is None or document is None:
+        return make_error(entry, "Auto document missing account or document")
+
+    return AutoDoc(start_date, end_date, account, document)
+
+
+def auto_document(entries: Entries, options_map, config_string=""):
+    del options_map, config_string # unused
+    errors = []
+    if len(entries) == 0:
+        return entries, errors
+
+    auto_docs: List[AutoDoc] = []
+    # keep track of balances that start on the same date, they are found first
+    bal_date = entries[0].date
+    balances = []
+
+    for entry in entries:
+        if entry.date > bal_date:
+            bal_date = entry.date
+            balances.clear()
+        if type(entry) is Custom and entry.type == "auto_document":
+            doc_data = custom_to_autodoc(entry)
+            if type(doc_data) is LoadError:
+                errors.append(doc_data)
+            elif type(doc_data) is AutoDoc:
+                auto_docs.append(doc_data)
+                for balance in balances:
+                    if balance.account == doc_data.account and 'document' not in balance.meta:
+                        balance.meta['document'] = doc_data.document
+
+        elif type(entry) is Transaction and 'document' not in entry.meta:
+            current_date = entry.date
+            for auto in auto_docs:
+                if auto.end_date < current_date:
+                    continue
+                match = False
+                for post in entry.postings:
+                    if post.account == auto.account:
+                        match = True
+                        break
+
+                if match:
+                    txd = entry._asdict()
+                    tagset = set()
+                    tagset.union(entry.tags)
+                    tagset.add("auto-doc")
+                    txd["tags"] = frozenset(tagset)
+                    entry.meta["document"] = auto.document
+                    break
+        elif type(entry) is Balance \
+                and 'document' not in entry.meta:
+            matched = False
+            for auto in auto_docs:
+                if auto.end_date >= entry.date and entry.account == auto.account:
+                    entry.meta["document"] = auto.document
+                    matched = True
+                    break
+
+            if not matched:
+                balances.append(entry)
+
+    return entries, errors
+