add ability to remove commands and commandsets

Issue #943

4 files changed, 148 insertions, 8 deletions

diff --git a/cmd2/cmd2.py b/cmd2/cmd2.py
index edf2a643..ef273d15 100644
--- a/cmd2/cmd2.py
+++ b/cmd2/cmd2.py
@@ -244,8 +244,8 @@ class Cmd(cmd.Cmd):
                                                 shortcuts=shortcuts)
 
         # Load modular commands
-        self._installed_functions: List[str] = []
-        self._installed_command_sets: List[CommandSet] = []
+        self._installed_functions = []  # type: List[str]
+        self._installed_command_sets = []  # type: List[CommandSet]
         if command_sets:
             for command_set in command_sets:
                 self.install_command_set(command_set)
@@ -469,7 +469,34 @@ class Cmd(cmd.Cmd):
                 delattr(self, attrib)
             raise
 
-    def install_command_function(self, cmd_name: str, cmd_func: Callable, cmd_completer: Callable, cmd_help: Callable):
+    def uninstall_command_set(self, cmdset: CommandSet):
+        """
+        Uninstalls an CommandSet and unloads all associated commands
+        :param cmdset: CommandSet to uninstall
+        """
+        if cmdset in self._installed_command_sets:
+            methods = inspect.getmembers(
+                cmdset,
+                predicate=lambda meth: inspect.ismethod(meth) and meth.__name__.startswith(COMMAND_FUNC_PREFIX))
+
+            for method in methods:
+                cmd_name = method[0][len(COMMAND_FUNC_PREFIX):]
+
+                delattr(self, COMMAND_FUNC_PREFIX + cmd_name)
+
+                if hasattr(self, COMPLETER_FUNC_PREFIX + cmd_name):
+                    delattr(self, COMPLETER_FUNC_PREFIX + cmd_name)
+                if hasattr(self, HELP_FUNC_PREFIX + cmd_name):
+                    delattr(self, HELP_FUNC_PREFIX + cmd_name)
+
+            cmdset.on_unregister(self)
+            self._installed_command_sets.remove(cmdset)
+
+    def install_command_function(self,
+                                 cmd_name: str,
+                                 cmd_func: Callable,
+                                 cmd_completer: Optional[Callable],
+                                 cmd_help: Optional[Callable]):
         """
         Installs a command by passing in functions for the command, completion, and help
 
@@ -483,7 +510,8 @@ class Cmd(cmd.Cmd):
         if not valid:
             raise ValueError("Invalid command name {!r}: {}".format(cmd_name, errmsg))
 
-        assert getattr(self, COMMAND_FUNC_PREFIX + cmd_name, None) is None, 'Duplicate command function registered: ' + cmd_name
+        assert getattr(self, COMMAND_FUNC_PREFIX + cmd_name, None) is None,\
+            'Duplicate command function registered: ' + cmd_name
         setattr(self, COMMAND_FUNC_PREFIX + cmd_name, types.MethodType(cmd_func, self))
         self._installed_functions.append(cmd_name)
         if cmd_completer is not None:
@@ -495,6 +523,20 @@ class Cmd(cmd.Cmd):
                 'Duplicate command help registered: ' + HELP_FUNC_PREFIX + cmd_name
             setattr(self, HELP_FUNC_PREFIX + cmd_name, types.MethodType(cmd_help, self))
 
+    def uninstall_command(self, cmd_name: str):
+        """
+        Uninstall an installed command and any associated completer or help functions
+        :param cmd_name: Command to uninstall
+        """
+        if cmd_name in self._installed_functions:
+            delattr(self, COMMAND_FUNC_PREFIX + cmd_name)
+
+            if hasattr(self, COMPLETER_FUNC_PREFIX + cmd_name):
+                delattr(self, COMPLETER_FUNC_PREFIX + cmd_name)
+            if hasattr(self, HELP_FUNC_PREFIX + cmd_name):
+                delattr(self, HELP_FUNC_PREFIX + cmd_name)
+            self._installed_functions.remove(cmd_name)
+
     def add_settable(self, settable: Settable) -> None:
         """
         Convenience method to add a settable parameter to ``self.settables``
diff --git a/cmd2/command_definition.py b/cmd2/command_definition.py
index a235525d..115cef64 100644
--- a/cmd2/command_definition.py
+++ b/cmd2/command_definition.py
@@ -87,6 +87,7 @@ def register_command(cmd_func: Callable[['Cmd', Union['Statement', 'argparse.Nam
             break
 
     _UNBOUND_COMMANDS.append((cmd_name, cmd_func, cmd_completer, cmd_help))
+    return cmd_func
 
 
 def with_default_category(category: str):
@@ -132,6 +133,12 @@ class CommandSet(object):
         to perform an initialization requiring access to the Cmd object.
 
         :param cmd: The cmd2 main application
-        :return: None
         """
         self._cmd = cmd
+
+    def on_unregister(self, cmd: 'Cmd'):
+        """
+        Called by ``cmd2.Cmd`` when a CommandSet is unregistered and removed.
+        :param cmd:
+        """
+        self._cmd = None
diff --git a/tests/conftest.py b/tests/conftest.py
index c07f7083..5b1a6f05 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -4,14 +4,13 @@ Cmd2 unit/functional testing
 """
 import sys
 from contextlib import redirect_stderr, redirect_stdout
-from typing import Dict, List, Optional, Union
+from typing import List, Optional, Union
 from unittest import mock
 
 from pytest import fixture
 
 import cmd2
 from cmd2.utils import StdSim
-from cmd2.constants import COMMAND_FUNC_PREFIX, CMD_ATTR_HELP_CATEGORY
 
 # Prefer statically linked gnureadline if available (for macOS compatibility due to issues with libedit)
 try:
diff --git a/tests/test_commandset.py b/tests/test_commandset.py
index acdb58b3..02fff7b2 100644
--- a/tests/test_commandset.py
+++ b/tests/test_commandset.py
@@ -31,8 +31,40 @@ def do_unbound(cmd: cmd2.Cmd, statement: cmd2.Statement):
     cmd.poutput('Unbound Command: {}'.format(statement.args))
 
 
+@cmd2.register_command
+@cmd2.with_category("AAA")
+def do_command_with_support(cmd: cmd2.Cmd, statement: cmd2.Statement):
+    """
+    This is an example of registering an unbound function
+
+    :param cmd:
+    :param statement:
+    :return:
+    """
+    cmd.poutput('Unbound Command: {}'.format(statement.args))
+
+
+def help_command_with_support(cmd: cmd2.Cmd):
+    cmd.poutput('Help for command_with_support')
+
+
+def complete_command_with_support(self, cmd: cmd2.Cmd, text: str, line: str, begidx: int, endidx: int) -> List[str]:
+    """Completion function for do_index_based"""
+    food_item_strs = ['Pizza', 'Ham', 'Ham Sandwich', 'Potato']
+    sport_item_strs = ['Bat', 'Basket', 'Basketball', 'Football', 'Space Ball']
+
+    index_dict = \
+        {
+            1: food_item_strs,  # Tab complete food items at index 1 in command line
+            2: sport_item_strs,  # Tab complete sport items at index 2 in command line
+            3: cmd.path_complete,  # Tab complete using path_complete function at index 3 in command line
+        }
+
+    return cmd.index_based_complete(text, line, begidx, endidx, index_dict=index_dict)
+
+
 @cmd2.with_default_category('Command Set')
-class TestCommandSet(cmd2.CommandSet):
+class CommandSetA(cmd2.CommandSet):
     def do_apple(self, cmd: cmd2.Cmd, statement: cmd2.Statement):
         cmd.poutput('Apple!')
 
@@ -69,6 +101,11 @@ def command_sets_app():
     app = WithCommandSets()
     return app
 
+@pytest.fixture()
+def command_sets_manual():
+    app = WithCommandSets(auto_load_commands=False)
+    return app
+
 
 def test_autoload_commands(command_sets_app):
     cmds_cats, cmds_doc, cmds_undoc, help_topics = command_sets_app._build_command_info()
@@ -83,4 +120,59 @@ def test_autoload_commands(command_sets_app):
     assert 'cranberry' in cmds_cats['Command Set']
 
 
+def test_load_commands(command_sets_manual):
+    cmds_cats, cmds_doc, cmds_undoc, help_topics = command_sets_manual._build_command_info()
+
+    assert 'AAA' not in cmds_cats
+
+    assert 'Alone' not in cmds_cats
+
+    assert 'Command Set' not in cmds_cats
+
+    command_sets_manual.install_command_function('unbound', do_unbound, None, None)
+
+    cmds_cats, cmds_doc, cmds_undoc, help_topics = command_sets_manual._build_command_info()
+
+    assert 'AAA' in cmds_cats
+    assert 'unbound' in cmds_cats['AAA']
+
+    assert 'Alone' not in cmds_cats
+
+    assert 'Command Set' not in cmds_cats
+
+    cmd_set = CommandSetA()
+
+    command_sets_manual.install_command_set(cmd_set)
+
+    cmds_cats, cmds_doc, cmds_undoc, help_topics = command_sets_manual._build_command_info()
+
+    assert 'AAA' in cmds_cats
+    assert 'unbound' in cmds_cats['AAA']
+
+    assert 'Alone' in cmds_cats
+    assert 'elderberry' in cmds_cats['Alone']
+
+    assert 'Command Set' in cmds_cats
+    assert 'cranberry' in cmds_cats['Command Set']
+
+    command_sets_manual.uninstall_command('unbound')
+
+    cmds_cats, cmds_doc, cmds_undoc, help_topics = command_sets_manual._build_command_info()
+
+    assert 'AAA' not in cmds_cats
+
+    assert 'Alone' in cmds_cats
+    assert 'elderberry' in cmds_cats['Alone']
+
+    assert 'Command Set' in cmds_cats
+    assert 'cranberry' in cmds_cats['Command Set']
+
+    command_sets_manual.uninstall_command_set(cmd_set)
+
+    cmds_cats, cmds_doc, cmds_undoc, help_topics = command_sets_manual._build_command_info()
+
+    assert 'AAA' not in cmds_cats
+
+    assert 'Alone' not in cmds_cats
 
+    assert 'Command Set' not in cmds_cats