From 3e4a284692edba7a87ed5caee641fb093a13d68a Mon Sep 17 00:00:00 2001
From: Felix Martin <mail@felixm.de>
Date: Sun, 2 Mar 2025 20:24:47 -0500
Subject: [PATCH] Add tags to mapping

---
 src/toldg/models.py  | 28 ++++++++++++++--------------
 src/toldg/predict.py |  4 +---
 src/toldg/process.py |  8 +-------
 src/toldg/utils.py   |  4 ++--
 src/toldg/write.py   | 39 ++++++++++++++++++++++++++-------------
 5 files changed, 44 insertions(+), 39 deletions(-)

diff --git a/src/toldg/models.py b/src/toldg/models.py
index ee4832d..2558a1f 100644
--- a/src/toldg/models.py
+++ b/src/toldg/models.py
@@ -52,6 +52,19 @@ class Config(BaseModel):
     categories: List[str]
 
 
+class Mapping(BaseModel):
+    """Class for transaction mapping from mappings file."""
+
+    class Config:
+        extra = "forbid"
+
+    account2: str
+    count: int = 1
+    narration: Optional[str] = None
+    payee: Optional[str] = None
+    tags: Optional[List[str]] = None
+
+
 class Transaction(BaseModel):
     """Class for ledger transaction to render into ldg file."""
 
@@ -67,17 +80,4 @@ class Transaction(BaseModel):
     description: str
     csv_file: str
     row: str
-    narration: Optional[str] = None
-    payee: Optional[str] = None
-
-
-class Mapping(BaseModel):
-    """Class for transaction mapping from mappings file."""
-
-    class Config:
-        extra = "forbid"
-
-    account2: str
-    count: int = 1
-    narration: Optional[str] = None
-    payee: Optional[str] = None
+    mapping: Optional[Mapping] = None
diff --git a/src/toldg/predict.py b/src/toldg/predict.py
index 0eca56c..59a56ec 100644
--- a/src/toldg/predict.py
+++ b/src/toldg/predict.py
@@ -28,9 +28,7 @@ def get_sort_categories():
 
 
 def add_account2(transactions: List[Transaction], categories: List[str]):
-    unmapped_transactions = list(
-        filter(lambda t: t.account2 == UNKNOWN_CATEGORY, transactions)
-    )
+    unmapped_transactions = list(filter(lambda t: t.mapping == None, transactions))
     if len(unmapped_transactions) == 0:
         return
     sort_categories = get_sort_categories()
diff --git a/src/toldg/process.py b/src/toldg/process.py
index a04802f..40b3fd0 100644
--- a/src/toldg/process.py
+++ b/src/toldg/process.py
@@ -74,13 +74,7 @@ def apply_mappings(transactions: List[Transaction], mappings: Dict[str, Mapping]
                 mapping.count > 0
             ), f"{mapping} used by {t} but count is not greater than '0'."
             mapping.count -= 1
-            t.account2 = mapping.account2
-
-            if mapping.narration:
-                t.narration = mapping.narration
-
-            if mapping.payee:
-                t.payee = mapping.payee
+            t.mapping = mapping
         else:
             logging.warning(f"No mapping for '{t}'.")
 
diff --git a/src/toldg/utils.py b/src/toldg/utils.py
index 07e30b4..0255377 100644
--- a/src/toldg/utils.py
+++ b/src/toldg/utils.py
@@ -81,9 +81,9 @@ def write_mappings(transactions: List[Transaction], mappings_file: Path):
             )
             mappings[t.row] = mapping
 
-    mappings = {k: v.dict() for k, v in mappings.items()}
+    mappings = {k: v.model_dump(exclude_none=True) for k, v in mappings.items()}
     with open(mappings_file, "w") as f:
-        json.dump(mappings, f, indent=4)
+        json.dump(mappings, f, indent=2)
 
 
 def read_mappings(mappings_file: Path) -> Dict[str, Mapping]:
diff --git a/src/toldg/write.py b/src/toldg/write.py
index e76e446..1a8bf85 100644
--- a/src/toldg/write.py
+++ b/src/toldg/write.py
@@ -5,7 +5,7 @@ from toldg.models import Config, Transaction
 from toldg.utils import category_to_bean
 
 BEANCOUNT_TRANSACTION_TEMPLATE = """
-{t.date} * {t.description}
+{t.date} * {description}{tags}
     {t.account2:<40}  {t.debit:<6} {t.currency}
     {t.account1:<40}  {t.credit:<6} {t.currency}
 """
@@ -13,28 +13,41 @@ BEANCOUNT_TRANSACTION_TEMPLATE = """
 
 def format(t):
     t.date = t.date.replace("/", "-")
-    if t.narration and t.payee:
-        # A transaction may have an optional “payee” and/or a “narration.”
-        t.description = f'"{t.payee}" "{t.narration}"'
-    elif t.narration:
-        # If you place a single string on a transaction line, it becomes its narration:
-        t.description = f'"{t.narration}"'
-    elif t.payee:
-        # If you want to set just a payee, put an empty narration string:
-        t.description = f'"{t.payee}" ""'
-    else:
-        t.description = f'"{t.description}"'
+
+    tags = ""
+    description = None
+    if t.mapping:
+        m = t.mapping
+        t.account2 = m.account2
+
+        if m.narration and m.payee:
+            # A transaction may have an optional “payee” and/or a “narration.”
+            description = f'"{m.payee}" "{m.narration}"'
+        elif m.narration:
+            # If you place a single string on a transaction line, it becomes its narration:
+            description = f'"{m.narration}"'
+        elif m.payee:
+            # If you want to set just a payee, put an empty narration string:
+            description = f'"{m.payee}" ""'
+        if m.tags:
+            tags = " #" + " #".join(m.tags)
+
+    if description is None:
+        description = f'"{t.description}"'
 
     if not t.debit.startswith("-"):
         t.debit = " " + t.debit
     if not t.credit.startswith("-"):
         t.credit = " " + t.credit
+
     t.account1 = category_to_bean(t.account1)
     t.account2 = category_to_bean(t.account2)
     if t.currency == "EUR":
         t.debit = t.debit.replace(".", "|").replace(",", ".").replace("|", ",")
         t.credit = t.credit.replace(".", "|").replace(",", ".").replace("|", ",")
-    return BEANCOUNT_TRANSACTION_TEMPLATE.format(t=t)
+    return BEANCOUNT_TRANSACTION_TEMPLATE.format(
+        t=t, description=description, tags=tags
+    )
 
 
 def render_to_file(transactions: List[Transaction], config: Config):