Added CommandSet.on_unregistered()

4 files changed, 33 insertions, 14 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index a22bfdf8..fe9c1ca6 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,7 +1,12 @@
-## 1.3.9 (September 01, 2020)
+## 1.3.9 (September 03, 2020)
+* Breaking Changes
+    * `CommandSet.on_unregister()` is now called as first step in unregistering a `CommandSet` and not
+    the last. `CommandSet.on_unregistered()` is now the last step.
 * Enhancements
-    * Added `on_registered()` callback to `CommandSet` class. This is called by `cmd2.Cmd` after a
-    `CommandSet` is registered and all its commands have been added to the CLI.
+    * Added `CommandSet.on_registered()`. This is called by `cmd2.Cmd` after a `CommandSet` is registered
+    and all its commands have been added to the CLI.
+    * Added `CommandSet.on_unregistered()`. This is called by `cmd2.Cmd` after a `CommandSet` is unregistered
+    and all its commands have been removed from the CLI.
 
 ## 1.3.8 (August 28, 2020)
 * Bug Fixes
diff --git a/cmd2/cmd2.py b/cmd2/cmd2.py
index 8fd3d243..f3a2d88d 100644
--- a/cmd2/cmd2.py
+++ b/cmd2/cmd2.py
@@ -505,6 +505,7 @@ class Cmd(cmd.Cmd):
             self._register_subcommands(cmdset)
             cmdset.on_registered()
         except Exception:
+            cmdset.on_unregister()
             for attrib in installed_attributes:
                 delattr(self, attrib)
             if cmdset in self._installed_command_sets:
@@ -512,7 +513,7 @@ class Cmd(cmd.Cmd):
             if cmdset in self._cmd_to_command_sets.values():
                 self._cmd_to_command_sets = \
                     {key: val for key, val in self._cmd_to_command_sets.items() if val is not cmdset}
-            cmdset.on_unregister()
+            cmdset.on_unregistered()
             raise
 
     def _install_command_function(self, command: str, command_wrapper: Callable, context=''):
@@ -560,6 +561,7 @@ class Cmd(cmd.Cmd):
         """
         if cmdset in self._installed_command_sets:
             self._check_uninstallable(cmdset)
+            cmdset.on_unregister()
             self._unregister_subcommands(cmdset)
 
             methods = inspect.getmembers(
@@ -585,7 +587,7 @@ class Cmd(cmd.Cmd):
                 if hasattr(self, HELP_FUNC_PREFIX + cmd_name):
                     delattr(self, HELP_FUNC_PREFIX + cmd_name)
 
-            cmdset.on_unregister()
+            cmdset.on_unregistered()
             self._installed_command_sets.remove(cmdset)
 
     def _check_uninstallable(self, cmdset: CommandSet):
diff --git a/cmd2/command_definition.py b/cmd2/command_definition.py
index 3c663054..64adaada 100644
--- a/cmd2/command_definition.py
+++ b/cmd2/command_definition.py
@@ -55,8 +55,9 @@ class CommandSet(object):
 
     def on_register(self, cmd) -> None:
         """
-        Called by cmd2.Cmd when a CommandSet is registered. Subclasses can override this
-        to perform an initialization requiring access to the Cmd object.
+        Called by cmd2.Cmd as the first step to registering a CommandSet. The commands defined in this class have
+        not be added to the CLI object at this point. Subclasses can override this to perform any initialization
+        requiring access to the Cmd object (e.g. configure commands and their parsers based on CLI state data).
 
         :param cmd: The cmd2 main application
         :type cmd: cmd2.Cmd
@@ -68,16 +69,22 @@ class CommandSet(object):
 
     def on_registered(self) -> None:
         """
-        Called by cmd2.Cmd after a CommandSet is registered and all its commands have been added
-        to the CLI. Subclasses can override this to perform custom steps.
+        Called by cmd2.Cmd after a CommandSet is registered and all its commands have been added to the CLI.
+        Subclasses can override this to perform custom steps related to the newly added commands (e.g. setting
+        them to a disabled state).
         """
         pass
 
     def on_unregister(self) -> None:
         """
-        Called by ``cmd2.Cmd`` when a CommandSet is unregistered and removed.
+        Called by ``cmd2.Cmd`` as the first step to unregistering a CommandSet. Subclasses can override this to
+        perform any cleanup steps which require their commands being registered in the CLI.
+        """
+        pass
 
-        :param cmd:
-        :type cmd: cmd2.Cmd
+    def on_unregistered(self) -> None:
+        """
+        Called by ``cmd2.Cmd`` after a CommandSet has been unregistered and all its commands removed from the CLI.
+        Subclasses can override this to perform remaining cleanup steps.
         """
         self._cmd = None
diff --git a/tests_isolated/test_commandset/test_commandset.py b/tests_isolated/test_commandset/test_commandset.py
index cbac7654..1685accf 100644
--- a/tests_isolated/test_commandset/test_commandset.py
+++ b/tests_isolated/test_commandset/test_commandset.py
@@ -30,8 +30,12 @@ class CommandSetA(CommandSetBase):
         print("in on_registered now")
 
     def on_unregister(self) -> None:
-        print("in on_unregister now")
         super().on_unregister()
+        print("in on_unregister now")
+
+    def on_unregistered(self) -> None:
+        super().on_unregistered()
+        print("in on_unregistered now")
 
     def do_apple(self, statement: cmd2.Statement):
         self._cmd.poutput('Apple!')
@@ -209,9 +213,10 @@ def test_load_commands(command_sets_manual, capsys):
     assert 'Alone' not in cmds_cats
     assert 'Fruits' not in cmds_cats
 
-    # Make sure unregistration callback ran
+    # Make sure unregistration callbacks ran
     out, err = capsys.readouterr()
     assert "in on_unregister now" in out
+    assert "in on_unregistered now" in out
 
     # uninstall a second time and verify no errors happen
     command_sets_manual.unregister_command_set(cmd_set)