Introduce new command error handling system

bennett-nguyen · Sep 14, 2024 · 2a0abd6 · 2a0abd6
1 parent 658a620
commit 2a0abd6
Show file tree

Hide file tree

Showing 22 changed files with 216 additions and 96 deletions.
diff --git a/src/base_command.py b/src/base_command.py
@@ -1,7 +1,8 @@
 from typing import Optional
-import argparse
 
 from src.app_state.app_state import AppState
+from src.override import ArgumentParserNoExit
+from src.exceptions import CommandException, ArgumentError
 
 class BaseCommand:
     """Represents a command with a parser for handling arguments.
@@ -15,8 +16,8 @@ class BaseCommand:
 
     def __init__(self, name: Optional[str] = None, usage: Optional[str] = None, description: Optional[str] = None, \
         epilog: Optional[str] = None, exit_on_error: bool = False):
-        self.parser = argparse.ArgumentParser(name, usage, description, epilog,
+        self.parser = ArgumentParserNoExit(name, usage, description, epilog,
                                             exit_on_error=exit_on_error)
 
-    def execute(self, args: list[str], app_state: AppState):
+    def execute(self, args: list[str], app_state: AppState) -> Optional[ArgumentError | CommandException]:
         ...
diff --git a/src/exceptions/__init__.py b/src/exceptions/__init__.py
@@ -0,0 +1,2 @@
+from argparse import ArgumentError
+from src.exceptions.command_exception import CommandException
diff --git a/src/exceptions/command_exception.py b/src/exceptions/command_exception.py
@@ -0,0 +1,10 @@
+class CommandException(Exception):
+    """
+    Custom exception class for handling command-related errors. This exception
+    allows for the inclusion of specific errors that can't be catched with
+    ArgumentError and is used in ArgumentParserNoExit.  
+    """
+
+    def __init__(self, message: str) -> None:
+        self.message = message
+        super().__init__(self.message)
diff --git a/src/exports/commands/config_cmd/list_queryfn.py b/src/exports/commands/config_cmd/list_queryfn.py
@@ -1,5 +1,8 @@
+from typing import Optional
+
 from src.app_state.app_state import AppState
 from src.base_command import BaseCommand
+from src.exceptions import ArgumentError, CommandException
 
 class ListQueryFN(BaseCommand):
     def __init__(self):
@@ -8,7 +11,7 @@ def __init__(self):
             description="List out available query functions.",
         )
 
-    def execute(self, args: list[str], app_state: AppState):
+    def execute(self, args: list[str], app_state: AppState)  -> Optional[ArgumentError | CommandException]:
         for fn_name in app_state.tree_manager.available_functions.keys():
             print(fn_name)
 

diff --git a/src/exports/commands/config_cmd/list_theme.py b/src/exports/commands/config_cmd/list_theme.py
@@ -1,5 +1,8 @@
+from typing import Optional
+
 from src.app_state.app_state import AppState
 from src.base_command import BaseCommand
+from src.exceptions import ArgumentError, CommandException
 
 class ListTheme(BaseCommand):
     def __init__(self):
@@ -8,7 +11,7 @@ def __init__(self):
             description="List out available themes.",
         )
 
-    def execute(self, args: list[str], app_state: AppState):
+    def execute(self, args: list[str], app_state: AppState) -> Optional[ArgumentError | CommandException]:
         for theme in app_state.theme_manager.available_themes:
             print(theme)
 

diff --git a/src/exports/commands/config_cmd/query_fn.py b/src/exports/commands/config_cmd/query_fn.py
@@ -1,6 +1,9 @@
 import argparse
+from typing import Optional
+
 from src.app_state.app_state import AppState
 from src.base_command import BaseCommand
+from src.exceptions import ArgumentError, CommandException
 
 class QueryFN(BaseCommand):
     def __init__(self):
@@ -11,15 +14,20 @@ def __init__(self):
 
         self.parser.add_argument("fn_name", type=str)
 
-    def execute(self, args: list[str], app_state: AppState):
-        parsed_args: argparse.Namespace = self.parser.parse_args(args)
-        tree_manager = app_state.tree_manager
-        segment_tree = app_state.tree_manager.segment_tree
+    def execute(self, args: list[str], app_state: AppState) -> Optional[ArgumentError | CommandException]:
+        try:
+            parsed_args: argparse.Namespace = self.parser.parse_args(args)
+            tree_manager = app_state.tree_manager
+            segment_tree = app_state.tree_manager.segment_tree
 
-        tree_manager.switch_function(parsed_args.fn_name)
-        segment_tree.rebuild()
-        tree_manager.generate_node_position()
-        tree_manager.compute_transformed_coordinates()
-        tree_manager.center_tree()
+            tree_manager.switch_function(parsed_args.fn_name)
+            segment_tree.rebuild()
+            tree_manager.generate_node_position()
+            tree_manager.compute_transformed_coordinates()
+            tree_manager.center_tree()
+        except argparse.ArgumentError as e:
+            return e
+        except CommandException as e:
+            return e
 
 query_fn_cmd = QueryFN()
diff --git a/src/exports/commands/config_cmd/theme.py b/src/exports/commands/config_cmd/theme.py
@@ -1,6 +1,9 @@
 import argparse
+from typing import Optional
+
 from src.app_state.app_state import AppState
 from src.base_command import BaseCommand
+from src.exceptions import ArgumentError, CommandException
 
 class Theme(BaseCommand):
     def __init__(self):
@@ -11,9 +14,14 @@ def __init__(self):
 
         self.parser.add_argument("theme_name", type=str)
 
-    def execute(self, args: list[str], app_state: AppState):
-        parsed_args: argparse.Namespace = self.parser.parse_args(args)
-        app_state.theme_manager.set_theme(parsed_args.theme_name)
-        app_state.cmdline_interface.set_theme(app_state.theme_manager.current_theme)
+    def execute(self, args: list[str], app_state: AppState) -> Optional[ArgumentError | CommandException]:
+        try:
+            parsed_args: argparse.Namespace = self.parser.parse_args(args)
+            app_state.theme_manager.set_theme(parsed_args.theme_name)
+            app_state.cmdline_interface.set_theme(app_state.theme_manager.current_theme)
+        except ArgumentError as e:
+            return e
+        except CommandException as e:
+            return e
 
 theme_cmd = Theme()
diff --git a/src/exports/commands/config_cmd/view_array.py b/src/exports/commands/config_cmd/view_array.py
@@ -1,6 +1,9 @@
+from typing import Optional
+
 from src.utils import VisibilityEnum
-from src.app_state.app_state import AppState
 from src.base_command import BaseCommand
+from src.app_state.app_state import AppState
+from src.exceptions import ArgumentError, CommandException
 
 class ViewArray(BaseCommand):
     def __init__(self):
@@ -9,7 +12,7 @@ def __init__(self):
             description="Toggle the visibility of the segment tree's array."
         )
 
-    def execute(self, args: list[str], app_state: AppState):
+    def execute(self, args: list[str], app_state: AppState) -> Optional[ArgumentError | CommandException]:
         app_state.rendering.visibility_dict[VisibilityEnum.ARRAY_FIELD] = not app_state.rendering.visibility_dict[VisibilityEnum.ARRAY_FIELD]
 
 view_array_cmd = ViewArray()
diff --git a/src/exports/commands/config_cmd/view_node_data.py b/src/exports/commands/config_cmd/view_node_data.py
@@ -1,6 +1,9 @@
+from typing import Optional
+
 from src.utils import VisibilityEnum
-from src.app_state.app_state import AppState
 from src.base_command import BaseCommand
+from src.app_state.app_state import AppState
+from src.exceptions import ArgumentError, CommandException
 
 class ViewNodeData(BaseCommand):
     def __init__(self):
@@ -9,7 +12,7 @@ def __init__(self):
             description="Toggle the visibility of a hovered node data."
         )
 
-    def execute(self, args: list[str], app_state: AppState):
+    def execute(self, args: list[str], app_state: AppState)  -> Optional[ArgumentError | CommandException]:
         app_state.rendering.visibility_dict[VisibilityEnum.NODE_DATA_FIELD] = not app_state.rendering.visibility_dict[VisibilityEnum.NODE_DATA_FIELD]
 
 view_node_data_cmd = ViewNodeData()
diff --git a/src/exports/commands/config_cmd/view_node_info.py b/src/exports/commands/config_cmd/view_node_info.py
@@ -1,7 +1,9 @@
-import argparse
+from typing import Optional
+
 from src.utils import VisibilityEnum
-from src.app_state.app_state import AppState
 from src.base_command import BaseCommand
+from src.app_state.app_state import AppState
+from src.exceptions import ArgumentError, CommandException
 
 class ViewNodeInfo(BaseCommand):
     def __init__(self):
@@ -10,7 +12,7 @@ def __init__(self):
             description="Toggle the visibility of a hovered node info."
         )
 
-    def execute(self, args: list[str], app_state: AppState):
+    def execute(self, args: list[str], app_state: AppState) -> Optional[ArgumentError | CommandException]:
         app_state.rendering.visibility_dict[VisibilityEnum.NODE_INFO_FIELD] = not app_state.rendering.visibility_dict[VisibilityEnum.NODE_INFO_FIELD]
 
 view_node_info_cmd = ViewNodeInfo()
diff --git a/src/exports/commands/rendering_cmd/highlight_range.py b/src/exports/commands/rendering_cmd/highlight_range.py
@@ -1,23 +1,30 @@
 import argparse
+from typing import Optional
+
 from src.utils import CommandRequestFields
 from src.app_state.app_state import AppState
 from src.base_command import BaseCommand
+from src.exceptions import ArgumentError, CommandException
 
 class HighlightRange(BaseCommand):
     def __init__(self):
         super().__init__(
             name="highlight-range",
             description="Highlight the nodes whose segments is within the specified range.",
         )
-        
+
         self.parser.add_argument("low", nargs="?", type=int, default=-1)
         self.parser.add_argument("high", nargs="?", type=int, default=-1)
-
-    def execute(self, args: list[str], app_state: AppState):
-        parsed_args: argparse.Namespace = self.parser.parse_args(args)
 
-        command_request_data = app_state.rendering.command_request_data
-        command_request_data[CommandRequestFields.HIGHLIGHT_RANGE_LOW] = parsed_args.low
-        command_request_data[CommandRequestFields.HIGHLIGHT_RANGE_HIGH] = parsed_args.high
+    def execute(self, args: list[str], app_state: AppState) -> Optional[ArgumentError | CommandException]:
+        try:
+            parsed_args: argparse.Namespace = self.parser.parse_args(args)
+            command_request_data = app_state.rendering.command_request_data
+            command_request_data[CommandRequestFields.HIGHLIGHT_RANGE_LOW] = parsed_args.low
+            command_request_data[CommandRequestFields.HIGHLIGHT_RANGE_HIGH] = parsed_args.high
+        except ArgumentError as e:
+            return e
+        except CommandException as e:
+            return e
 
 highlight_range_cmd = HighlightRange()
diff --git a/src/exports/commands/tree_cmd/clear.py b/src/exports/commands/tree_cmd/clear.py
@@ -1,6 +1,8 @@
-import argparse
-from src.app_state.app_state import AppState
+from typing import Optional
+
 from src.base_command import BaseCommand
+from src.app_state.app_state import AppState
+from src.exceptions import ArgumentError, CommandException
 
 class Clear(BaseCommand):
     def __init__(self):
@@ -9,7 +11,7 @@ def __init__(self):
             description="Clear the segment tree's array.",
         )
 
-    def execute(self, args: list[str], app_state: AppState):
+    def execute(self, args: list[str], app_state: AppState) -> Optional[ArgumentError | CommandException]:
         tree_manager = app_state.tree_manager
         tree_manager.segment_tree.array.clear()
         tree_manager.generate_node_position()

diff --git a/src/exports/commands/tree_cmd/extend.py b/src/exports/commands/tree_cmd/extend.py
@@ -1,6 +1,9 @@
 import argparse
-from src.app_state.app_state import AppState
+from typing import Optional
+
 from src.base_command import BaseCommand
+from src.app_state.app_state import AppState
+from src.exceptions import ArgumentError, CommandException
 
 class Extend(BaseCommand):
     def __init__(self):
@@ -12,19 +15,24 @@ def __init__(self):
         self.parser.add_argument("sequence", type=int, nargs="+")
         self.parser.add_argument("-i", "--index", "-index", type=int, nargs="?", default=-1)
 
-    def execute(self, args: list[str], app_state: AppState):
-        parsed_args: argparse.Namespace = self.parser.parse_args(args)
-        tree_manager = app_state.tree_manager
-        segment_tree = tree_manager.segment_tree
+    def execute(self, args: list[str], app_state: AppState) -> Optional[ArgumentError | CommandException]:
+        try:
+            parsed_args: argparse.Namespace = self.parser.parse_args(args)
+            tree_manager = app_state.tree_manager
+            segment_tree = tree_manager.segment_tree
 
-        index_to_extend = len(segment_tree.array)
-        if (parsed_args.index != -1):
-            index_to_extend = parsed_args.index
+            index_to_extend = len(segment_tree.array)
+            if (parsed_args.index != -1):
+                index_to_extend = parsed_args.index
 
-        segment_tree.array[index_to_extend:index_to_extend] = parsed_args.sequence
-        segment_tree.rebuild()
-        tree_manager.generate_node_position()
-        tree_manager.compute_transformed_coordinates()
-        tree_manager.center_tree()
+            segment_tree.array[index_to_extend:index_to_extend] = parsed_args.sequence
+            segment_tree.rebuild()
+            tree_manager.generate_node_position()
+            tree_manager.compute_transformed_coordinates()
+            tree_manager.center_tree()
+        except ArgumentError as e:
+            return e
+        except CommandException as e:
+            return e
 
 extend_cmd = Extend()
diff --git a/src/exports/commands/tree_cmd/home.py b/src/exports/commands/tree_cmd/home.py
@@ -1,8 +1,8 @@
-from src.app_state.app_state import AppState
+from typing import Optional
 
 from src.base_command import BaseCommand
-from src.window import pygame_window
-from src.utils import const
+from src.app_state.app_state import AppState
+from src.exceptions import ArgumentError, CommandException
 
 class Home(BaseCommand):
     def __init__(self):
@@ -11,7 +11,7 @@ def __init__(self):
             description="Move the tree to its original position."
         )
 
-    def execute(self, args: list[str], app_state: AppState):
+    def execute(self, args: list[str], app_state: AppState) -> Optional[ArgumentError | CommandException]:
         app_state.tree_manager.center_tree()
 
 home_cmd = Home()
diff --git a/src/exports/commands/tree_cmd/insert.py b/src/exports/commands/tree_cmd/insert.py
@@ -1,6 +1,9 @@
 import argparse
-from src.app_state.app_state import AppState
+from typing import Optional
+
 from src.base_command import BaseCommand
+from src.app_state.app_state import AppState
+from src.exceptions import ArgumentError, CommandException
 
 class Insert(BaseCommand):
     def __init__(self):
@@ -12,19 +15,24 @@ def __init__(self):
         self.parser.add_argument("value", type=int)
         self.parser.add_argument("index", type=int, default=-1, nargs='?')
 
-    def execute(self, args: list[str], app_state: AppState):
-        parsed_args: argparse.Namespace = self.parser.parse_args(args)
-        tree_manager = app_state.tree_manager
-        segment_tree = tree_manager.segment_tree
-
-        index_to_insert = len(segment_tree.array)
-        if (parsed_args.index != -1):
-            index_to_insert = parsed_args.index
+    def execute(self, args: list[str], app_state: AppState) -> Optional[ArgumentError | CommandException]:
+        try:
+            parsed_args: argparse.Namespace = self.parser.parse_args(args)
+            tree_manager = app_state.tree_manager
+            segment_tree = tree_manager.segment_tree
+
+            index_to_insert = len(segment_tree.array)
+            if (parsed_args.index != -1):
+                index_to_insert = parsed_args.index
 
-        segment_tree.array.insert(index_to_insert, parsed_args.value)
-        segment_tree.rebuild()
-        tree_manager.generate_node_position()
-        tree_manager.compute_transformed_coordinates()
-        tree_manager.center_tree()
+            segment_tree.array.insert(index_to_insert, parsed_args.value)
+            segment_tree.rebuild()
+            tree_manager.generate_node_position()
+            tree_manager.compute_transformed_coordinates()
+            tree_manager.center_tree()
+        except ArgumentError as e:
+            return e
+        except CommandException as e:
+            return e
 
 insert_cmd = Insert()