From 6bafb3eae754e34ec9d411976d0c0fa09a91098d Mon Sep 17 00:00:00 2001
From: kraanzu <kraanzu@gmail.com>
Date: Wed, 18 Sep 2024 09:30:51 +0530
Subject: [PATCH] refactor: add decorator for requiring node highlight

---
 dooit/ui/widgets/trees/_decorators.py | 10 +++++++++
 dooit/ui/widgets/trees/base_tree.py   | 30 ++++++++++++++-------------
 dooit/ui/widgets/trees/model_tree.py  | 17 ++++-----------
 3 files changed, 30 insertions(+), 27 deletions(-)

diff --git a/dooit/ui/widgets/trees/_decorators.py b/dooit/ui/widgets/trees/_decorators.py
index 5e9d1c4e..d5a9b3e2 100644
--- a/dooit/ui/widgets/trees/_decorators.py
+++ b/dooit/ui/widgets/trees/_decorators.py
@@ -33,3 +33,13 @@ def wrapper(self: "ModelTree", *args, **kwargs) -> Any:
         self.force_refresh()
 
     return wrapper
+
+
+def require_highlighted_node(func: Callable) -> Callable:
+    def wrapper(self: "ModelTree", *args, **kwargs) -> Any:
+        if self.highlighted is None:
+            raise ValueError("No node is currently highlighted")
+
+        return func(self, *args, **kwargs)
+
+    return wrapper
diff --git a/dooit/ui/widgets/trees/base_tree.py b/dooit/ui/widgets/trees/base_tree.py
index 7b3cf6b5..c617de61 100644
--- a/dooit/ui/widgets/trees/base_tree.py
+++ b/dooit/ui/widgets/trees/base_tree.py
@@ -4,6 +4,7 @@
 from dooit.api.todo import Todo
 from dooit.api.workspace import Workspace
 from collections import defaultdict
+from dooit.ui.widgets.trees._decorators import require_highlighted_node
 
 ModelType = Union[Todo, Workspace]
 
@@ -24,20 +25,21 @@ def tui(self) -> "Dooit":
         raise ValueError("App is not a Dooit instance")
 
     @property
+    @require_highlighted_node
     def node(self) -> Option:
-        if self.highlighted is None:
-            raise ValueError("No node is currently highlighted")
-
+        assert self.highlighted is not None
         return self.get_option_at_index(self.highlighted)
 
-    def action_cursor_down(self) -> None:
-        if self.highlighted == len(self._options) - 1:
-            return
-
-        return super().action_cursor_down()
-
-    def action_cursor_up(self) -> None:
-        if self.highlighted == 0:
-            return
-
-        return super().action_cursor_up()
+    # TODO: Uncomment this:
+    #
+    # def action_cursor_down(self) -> None:
+    #     if self.highlighted == len(self._options) - 1:
+    #         return
+    #
+    #     return super().action_cursor_down()
+    #
+    # def action_cursor_up(self) -> None:
+    #     if self.highlighted == 0:
+    #         return
+    #
+    #     return super().action_cursor_up()
diff --git a/dooit/ui/widgets/trees/model_tree.py b/dooit/ui/widgets/trees/model_tree.py
index 1c4b78a6..7d0e217d 100644
--- a/dooit/ui/widgets/trees/model_tree.py
+++ b/dooit/ui/widgets/trees/model_tree.py
@@ -7,7 +7,7 @@
 from dooit.ui.widgets.renderers.base_renderer import BaseRenderer
 from .base_tree import BaseTree
 from ._render_dict import RenderDict
-from ._decorators import fix_highlight, refresh_tree
+from ._decorators import fix_highlight, refresh_tree, require_highlighted_node
 
 ModelType = TypeVar("ModelType", bound=Union[Todo, Workspace])
 RenderDictType = TypeVar("RenderDictType", bound=RenderDict)
@@ -145,10 +145,8 @@ def _toggle_expand_node(self, _id: str) -> None:
         else:
             self._expand_node(_id)
 
+    @require_highlighted_node
     def toggle_expand(self) -> None:
-        if self.highlighted is None or not self.node.id:
-            return
-
         self._toggle_expand_node(self.node.id)
 
     def _toggle_expand_parent(self, _id: str) -> None:
@@ -161,13 +159,8 @@ def _toggle_expand_parent(self, _id: str) -> None:
         self.highlight_id(parent_id)
         self._toggle_expand_node(parent_id)
 
+    @require_highlighted_node
     def toggle_expand_parent(self) -> None:
-        if self.highlighted is None:
-            return
-
-        if not self.node.id:
-            return
-
         self._toggle_expand_parent(self.node.id)
 
     def _create_child_node(self) -> ModelType:
@@ -201,10 +194,8 @@ def add_sibling(self):
         self.start_edit("description")
 
     @refresh_tree
+    @require_highlighted_node
     def remove_node(self):
-        if self.highlighted is None:
-            return
-
         self.current_model.drop()
 
     @refresh_tree