From 691ce30a68c51e3cd57e21d07b0bc0137048c043 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sergio=20Mart=C3=ADnez=20Portela?=
 <sergio@codigoparallevar.com>
Date: Mon, 7 Oct 2024 23:22:44 +0200
Subject: [PATCH] Simplify state setting, update `.is_todo`/`.is_done` props.

---
 org_rw/org_rw.py  | 42 ++++++++++++++++++++++++++---
 tests/test_org.py | 69 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 108 insertions(+), 3 deletions(-)

diff --git a/org_rw/org_rw.py b/org_rw/org_rw.py
index a466d4f..31b904c 100644
--- a/org_rw/org_rw.py
+++ b/org_rw/org_rw.py
@@ -337,7 +337,7 @@ class Headline:
         self.priority = priority
         self.title_start = title_start
         self.title = parse_content_block([RawLine(linenum=start_line, line=title)])
-        self.state = state
+        self._state = state
         self.tags_start = tags_start
         self.shallow_tags = tags
         self.contents = contents
@@ -726,6 +726,42 @@ class Headline:
     def id(self, value):
         self.set_property("ID", value)
 
+    @property
+    def state(self) -> HeadlineState:
+        return self._state
+
+    @state.setter
+    def state(self, new_state: Union[None, str, HeadlineState]) -> None:
+        """
+        Update the state of a Headline. If the state is a known one it will update it's TODO/DONE properties.
+
+        Args:
+            new_state (str|HeadlineState): New state, either it's literal value or it's structure.
+        """
+        if new_state is None:
+            self.is_todo = False
+            self.is_done = False
+            # TODO: Check & log if appropriate?
+            self._state = None
+            return
+
+        if isinstance(new_state, str):
+            new_state = HeadlineState(name=new_state)
+
+        state_name = new_state["name"]
+        if state_name in [kw["name"] for kw in self.doc.todo_keywords]:
+            self.is_todo = True
+            self.is_done = False
+            # TODO: Check & log if appropriate?
+        elif state_name in [kw["name"] for kw in self.doc.done_keywords]:
+            self.is_todo = False
+            self.is_done = True
+            # TODO: Check, log &  if appropriate?
+        else:
+            # TODO: Should we raise a warning, raise an exception, update the is_todo/is_done?
+            pass
+        self._state = new_state
+
     @property
     def clock(self):
         times = []
@@ -2378,8 +2414,8 @@ class OrgDoc:
             tags = ":" + ":".join(headline.shallow_tags) + ":"
 
         state = ""
-        if headline.state:
-            state = headline.state["name"] + " "
+        if headline._state:
+            state = headline._state["name"] + " "
 
         raw_title = token_list_to_raw(headline.title.contents)
         tags_padding = ""
diff --git a/tests/test_org.py b/tests/test_org.py
index f6b6be4..a1fdff1 100644
--- a/tests/test_org.py
+++ b/tests/test_org.py
@@ -955,6 +955,75 @@ class TestSerde(unittest.TestCase):
         h1_2_h2 = h1_2.children[0]
         self.assertEqual(sorted(h1_2_h2.tags), ["otherh2tag"])
 
+    def test_update_headline_from_none_to_todo(self):
+        orig = "* First entry"
+        doc = loads(orig)
+        self.assertEqual(doc.headlines[0].is_todo, False)
+        self.assertEqual(doc.headlines[0].is_done, False)
+        self.assertEqual(doc.headlines[0].state, None)
+
+        doc.headlines[0].state = "TODO"
+        self.assertEqual(doc.headlines[0].is_todo, True)
+        self.assertEqual(doc.headlines[0].is_done, False)
+        self.assertEqual(doc.headlines[0].state["name"], "TODO")
+
+        self.assertEqual(dumps(doc), "* TODO First entry")
+
+    def test_update_headline_from_none_to_done(self):
+        orig = "* First entry"
+        doc = loads(orig)
+        self.assertEqual(doc.headlines[0].is_todo, False)
+        self.assertEqual(doc.headlines[0].is_done, False)
+        self.assertEqual(doc.headlines[0].state, None)
+
+        doc.headlines[0].state = org_rw.HeadlineState(name="DONE")
+        self.assertEqual(doc.headlines[0].is_todo, False)
+        self.assertEqual(doc.headlines[0].is_done, True)
+        self.assertEqual(doc.headlines[0].state["name"], "DONE")
+
+        self.assertEqual(dumps(doc), "* DONE First entry")
+
+    def test_update_headline_from_todo_to_none(self):
+        orig = "* TODO First entry"
+        doc = loads(orig)
+        self.assertEqual(doc.headlines[0].is_todo, True)
+        self.assertEqual(doc.headlines[0].is_done, False)
+        self.assertEqual(doc.headlines[0].state["name"], "TODO")
+
+        doc.headlines[0].state = None
+        self.assertEqual(doc.headlines[0].is_todo, False)
+        self.assertEqual(doc.headlines[0].is_done, False)
+        self.assertEqual(doc.headlines[0].state, None)
+
+        self.assertEqual(dumps(doc), "* First entry")
+
+    def test_update_headline_from_todo_to_done(self):
+        orig = "* TODO First entry"
+        doc = loads(orig)
+        self.assertEqual(doc.headlines[0].is_todo, True)
+        self.assertEqual(doc.headlines[0].is_done, False)
+        self.assertEqual(doc.headlines[0].state["name"], "TODO")
+
+        doc.headlines[0].state = "DONE"
+        self.assertEqual(doc.headlines[0].is_todo, False)
+        self.assertEqual(doc.headlines[0].is_done, True)
+        self.assertEqual(doc.headlines[0].state["name"], "DONE")
+        self.assertEqual(dumps(doc), "* DONE First entry")
+
+    def test_update_headline_from_done_to_todo(self):
+        orig = "* DONE First entry"
+        doc = loads(orig)
+        self.assertEqual(doc.headlines[0].is_todo, False)
+        self.assertEqual(doc.headlines[0].is_done, True)
+        self.assertEqual(doc.headlines[0].state["name"], "DONE")
+
+        doc.headlines[0].state = org_rw.HeadlineState(name="TODO")
+        self.assertEqual(doc.headlines[0].is_todo, True)
+        self.assertEqual(doc.headlines[0].is_done, False)
+        self.assertEqual(doc.headlines[0].state["name"], "TODO")
+
+        self.assertEqual(dumps(doc), "* TODO First entry")
+
 
 def print_tree(tree, indentation=0, headline=None):
     for element in tree:
-- 
2.45.2