Merge pull request #990 from python-cmd2/on_registered

Added callbacks to CommandSet

2 files changed, 26 insertions, 9 deletions

diff --git a/cmd2/cmd2.py b/cmd2/cmd2.py
index d768085a..f3a2d88d 100644
--- a/cmd2/cmd2.py
+++ b/cmd2/cmd2.py
@@ -503,7 +503,9 @@ class Cmd(cmd.Cmd):
             self._installed_command_sets.append(cmdset)
 
             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:
@@ -511,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=''):
@@ -559,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(
@@ -584,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 27a044bc..64adaada 100644
--- a/cmd2/command_definition.py
+++ b/cmd2/command_definition.py
@@ -53,10 +53,11 @@ class CommandSet(object):
     def __init__(self):
         self._cmd = None  # type: Optional[cmd2.Cmd]
 
-    def on_register(self, cmd):
+    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
@@ -66,11 +67,24 @@ class CommandSet(object):
         else:
             raise CommandSetRegistrationError('This CommandSet has already been registered')
 
-    def on_unregister(self):
+    def on_registered(self) -> None:
         """
-        Called by ``cmd2.Cmd`` when a CommandSet is unregistered and removed.
+        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
 
-        :param cmd:
-        :type cmd: cmd2.Cmd
+    def on_unregister(self) -> None:
+        """
+        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
+
+    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