Added way to disable commands

1 files changed, 103 insertions, 4 deletions

diff --git a/cmd2/cmd2.py b/cmd2/cmd2.py
index b3a61212..8a5d31d3 100644
--- a/cmd2/cmd2.py
+++ b/cmd2/cmd2.py
@@ -279,6 +279,14 @@ class EmptyStatement(Exception):
     pass
 
 
+class DisabledCommand:
+    """Contains data about a disabled command"""
+    def __init__(self):
+        # These are used to restore the original functions when the command is enabled
+        self.command_function = None
+        self.help_function = None
+
+
 class Cmd(cmd.Cmd):
     """An easy but powerful framework for writing line-oriented command interpreters.
 
@@ -521,6 +529,11 @@ class Cmd(cmd.Cmd):
         # being printed by a command.
         self.terminal_lock = threading.RLock()
 
+        # Commands that have been disabled from use. This is to support commands that are only available
+        # during specific states of the application. This dictionary's keys are the command names and its
+        # values are DisabledCommand objects.
+        self.disabled_commands = dict()
+
     # -----  Methods related to presenting output to the user -----
 
     @property
@@ -1562,7 +1575,7 @@ class Cmd(cmd.Cmd):
                 if name.startswith(COMMAND_FUNC_PREFIX) and callable(getattr(self, name))]
 
     def get_visible_commands(self) -> List[str]:
-        """Returns a list of commands that have not been hidden."""
+        """Returns a list of commands that have not been hidden or disabled."""
         commands = self.get_all_commands()
 
         # Remove the hidden commands
@@ -1570,6 +1583,11 @@ class Cmd(cmd.Cmd):
             if name in commands:
                 commands.remove(name)
 
+        # Remove the disabled commands
+        for name in self.disabled_commands:
+            if name in commands:
+                commands.remove(name)
+
         return commands
 
     def get_alias_names(self) -> List[str]:
@@ -1953,7 +1971,7 @@ class Cmd(cmd.Cmd):
     def onecmd(self, statement: Union[Statement, str]) -> bool:
         """ This executes the actual do_* method for a command.
 
-        If the command provided doesn't exist, then it executes _default() instead.
+        If the command provided doesn't exist, then it executes default() instead.
 
         :param statement: intended to be a Statement instance parsed command from the input stream, alternative
                           acceptance of a str is present only for backward compatibility with cmd
@@ -3186,13 +3204,15 @@ class Cmd(cmd.Cmd):
 
         # -v must be used alone with no other options
         if args.verbose:
-            if args.clear or args.edit or args.output_file or args.run or args.transcript or args.expanded or args.script:
+            if (args.clear or args.edit or args.output_file or args.run or
+                    args.transcript or args.expanded or args.script):
                 self.poutput("-v can not be used with any other options")
                 self.poutput(self.history_parser.format_usage())
                 return
 
         # -s and -x can only be used if none of these options are present: [-c -r -e -o -t]
-        if (args.script or args.expanded) and (args.clear or args.edit or args.output_file or args.run or args.transcript):
+        if (args.script or args.expanded) and (args.clear or args.edit or args.output_file or args.run or
+                                               args.transcript):
             self.poutput("-s and -x can not be used with -c, -r, -e, -o, or -t")
             self.poutput(self.history_parser.format_usage())
             return
@@ -3598,6 +3618,85 @@ class Cmd(cmd.Cmd):
         else:
             raise RuntimeError("another thread holds terminal_lock")
 
+    def enable_command(self, command: str) -> None:
+        """
+        Enable a command by restoring its functions
+        :param command: the command being enabled
+        """
+        # If the commands is already enabled, then return
+        if command not in self.disabled_commands:
+            return
+
+        help_func_name = HELP_FUNC_PREFIX + command
+
+        # Restore the command and help functions to their original values
+        dc = self.disabled_commands[command]
+        setattr(self, self.cmd_func_name(command), dc.command_function)
+
+        if dc.help_function is None:
+            delattr(self, help_func_name)
+        else:
+            setattr(self, help_func_name, dc.help_function)
+
+        # Remove the disabled command entry
+        del self.disabled_commands[command]
+
+    def enable_category(self, category: str) -> None:
+        """
+        Enable an entire category of commands
+        :param category: the category to enable
+        """
+        for cmd_name in self.disabled_commands:
+            dc = self.disabled_commands[cmd_name]
+            cmd_category = getattr(dc.command_function, HELP_CATEGORY, None)
+            if cmd_category is not None and cmd_category == category:
+                self.enable_command(cmd_name)
+
+    def disable_command(self, command: str, message_to_print: str) -> None:
+        """
+        Disable a command and overwrite its functions
+        :param command: the command being disabled
+        :param message_to_print: what to print when this command or its help function is run while disabled
+        """
+        import functools
+
+        # If the commands is already disabled, then return
+        if command in self.disabled_commands:
+            return
+
+        # Make sure this is an actual command
+        command_function = self.cmd_func(command)
+        if command_function is None:
+            raise AttributeError("{} does not refer to a command".format(command))
+
+        help_func_name = HELP_FUNC_PREFIX + command
+
+        # Add the disabled command record
+        dc = DisabledCommand()
+        dc.command_function = command_function
+        dc.help_function = getattr(self, help_func_name, None)
+        self.disabled_commands[command] = dc
+
+        # Overwrite the command and help functions to print the message
+        setattr(self, self.cmd_func_name(command), functools.partial(self.poutput, message_to_print + '\n'))
+        setattr(self, help_func_name, functools.partial(self.poutput, message_to_print + '\n'))
+
+    def disable_category(self, category: str, message_to_print: str) -> None:
+        """
+        Disable an entire category of commands
+        :param category: the category to disable
+        :param message_to_print: what to print when anything in this category is run while disabled
+        """
+        all_commands = self.get_all_commands()
+
+        for cmd_name in all_commands:
+            func = self.cmd_func(cmd_name)
+            cmd_category = getattr(func, HELP_CATEGORY, None)
+
+            # If this command is in the category, then disable it
+            if cmd_category is not None and cmd_category == category:
+                self.disable_command(cmd_name, message_to_print)
+
     def cmdloop(self, intro: Optional[str] = None) -> None:
         """This is an outer wrapper around _cmdloop() which deals with extra features provided by cmd2.