From 95fea80f9086581f20fdb692cafbd3f98d3423b0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sergio=20Mart=C3=ADnez=20Portela?=
 <sergio@codigoparallevar.com>
Date: Sat, 20 Jul 2024 17:29:25 +0200
Subject: [PATCH] Refactor headline state parsing.

- Add separate function to parse states.
- Handle edge case when no `|` is used to split TODO and DONE states.
- Add typing to the states to future-proof for handling keyboard shortcuts and actions on state changes.
---
 org_rw/org_rw.py | 55 ++++++++++++++++++++++++++++++++++++++----------
 1 file changed, 44 insertions(+), 11 deletions(-)

diff --git a/org_rw/org_rw.py b/org_rw/org_rw.py
index 3a775a3..1f0df7a 100644
--- a/org_rw/org_rw.py
+++ b/org_rw/org_rw.py
@@ -9,7 +9,7 @@ import re
 import sys
 from datetime import date, datetime, timedelta
 from enum import Enum
-from typing import cast, Iterator, List, Literal, Optional, Tuple, Union
+from typing import cast, Iterator, List, Literal, Optional, Tuple, TypedDict, Union
 
 from .types import HeadlineDict
 
@@ -107,6 +107,15 @@ CodeSnippet = collections.namedtuple("CodeSnippet", ("name", "content", "result"
 NON_FINISHED_GROUPS = (type(None), dom.ListGroupNode, dom.ResultsDrawerNode, dom.PropertyDrawerNode)
 FREE_GROUPS = (dom.CodeBlock,)
 
+# States
+class HeadlineState(TypedDict):
+    # To be extended to handle keyboard shortcuts
+    name: str
+
+class OrgDocDeclaredStates(TypedDict):
+    not_completed: List[HeadlineState]
+    completed: List[HeadlineState]
+
 
 class NonReproducibleDocument(Exception):
     """
@@ -1863,29 +1872,53 @@ def dump_delimiters(line: DelimiterLine):
     return (line.linenum, line.line)
 
 
+def parse_todo_done_keywords(line: str) -> OrgDocDeclaredStates:
+    clean_line = re.sub(r"\([^)]+\)", "", line)
+    if '|' in clean_line:
+        todo_kws, done_kws = clean_line.split("|", 1)
+        has_split = True
+    else:
+        # Standard behavior in this case is: the last state is the one considered as DONE
+        todo_kws = clean_line
+
+    todo_keywords = re.sub(r"\s{2,}", " ", todo_kws.strip()).split()
+    if has_split:
+        done_keywords = re.sub(r"\s{2,}", " ", done_kws.strip()).split()
+    else:
+        done_keywods = [todo_keywords[-1]]
+        todo_keywords = todo_keywords[:-1]
+
+    return {
+        "not_completed": [
+            HeadlineState(name=keyword)
+            for keyword in todo_keywords
+        ],
+        "completed": [
+            HeadlineState(name=keyword)
+            for keyword in done_keywords
+        ],
+    }
+
+
 class OrgDoc:
     def __init__(
         self, headlines, keywords, contents, list_items, structural, properties,
         environment=BASE_ENVIRONMENT,
     ):
-        self.todo_keywords = DEFAULT_TODO_KEYWORDS
-        self.done_keywords = DEFAULT_DONE_KEYWORDS
+        self.todo_keywords = [HeadlineState(name=kw) for kw in DEFAULT_TODO_KEYWORDS]
+        self.done_keywords = [HeadlineState(name=kw) for kw in DEFAULT_DONE_KEYWORDS]
 
         keywords_set_in_file = False
         for keyword in keywords:
             if keyword.key in ("TODO", "SEQ_TODO"):
-                todo_kws, done_kws = re.sub(r"\([^)]+\)", "", keyword.value).split("|", 1)
-
-                self.todo_keywords = re.sub(r"\s{2,}", " ", todo_kws.strip()).split()
-                self.done_keywords = re.sub(r"\s{2,}", " ", done_kws.strip()).split()
+                states = parse_todo_done_keywords(keyword.value)
+                self.todo_keywords, self.done_keywords = states['not_completed'], states['completed']
                 keywords_set_in_file = True
 
         if not keywords_set_in_file and 'org-todo-keywords' in environment:
             # Read keywords from environment
-            todo_kws, done_kws = re.sub(r"\([^)]+\)", "", environment['org-todo-keywords']).split("|", 1)
-
-            self.todo_keywords = re.sub(r"\s{2,}", " ", todo_kws.strip()).split()
-            self.done_keywords = re.sub(r"\s{2,}", " ", done_kws.strip()).split()
+            states = parse_todo_done_keywords(environment['org-todo-keywords'])
+            self.todo_keywords, self.done_keywords = states['not_completed'], states['completed']
 
         self.keywords: List[Property] = keywords
         self.contents: List[RawLine] = contents