Add command finalization hooks

4 files changed, 192 insertions, 9 deletions

diff --git a/cmd2/cmd2.py b/cmd2/cmd2.py
index 0bb7509f..140d7034 100644
--- a/cmd2/cmd2.py
+++ b/cmd2/cmd2.py
@@ -1687,9 +1687,19 @@ class Cmd(cmd.Cmd):
         :return: True if cmdloop() should exit, False otherwise
         """
         import datetime
+
         stop = False
         try:
             statement = self._complete_statement(line)
+        except EmptyStatement:
+            return self._run_cmdfinalization_hooks(stop, None)
+        except ValueError as ex:
+            # If shlex.split failed on syntax, let user know whats going on
+            self.perror("Invalid syntax: {}".format(ex), traceback_war=False)
+            return stop
+
+        # now that we have a statement, run it with all the hooks
+        try:
             # call the postparsing hooks
             data = plugin.PostparsingData(False, statement)
             for func in self._postparsing_hooks:
@@ -1740,15 +1750,26 @@ class Cmd(cmd.Cmd):
                 if self.allow_redirection and self.redirecting:
                     self._restore_output(statement)
         except EmptyStatement:
+            # don't do anything, but do allow command finalization hooks to run
             pass
-        except ValueError as ex:
-            # TODO move this except to way above, so ValueError is only caught for shlex errors
-            # If shlex.split failed on syntax, let user know whats going on
-            self.perror("Invalid syntax: {}".format(ex), traceback_war=False)
         except Exception as ex:
             self.perror(ex)
         finally:
+            return self._run_cmdfinalization_hooks(stop, statement)
+
+    def _run_cmdfinalization_hooks(self, stop: bool, statement: Statement) -> bool:
+        """Run the command finalization hooks"""
+        try:
+            data = plugin.CommandFinalizationData(stop, statement)
+            for func in self._cmdfinalization_hooks:
+                data = func(data)
+            # retrieve the final value of stop, ignoring any
+            # modifications to the statement
+            stop = data.stop
+            # postparsing_postcmd is deprecated
             return self.postparsing_postcmd(stop)
+        except Exception as ex:
+            self.perror(ex)
 
     def runcmds_plus_hooks(self, cmds: List[str]) -> bool:
         """Convenience method to run multiple commands by onecmd_plus_hooks.
@@ -3138,9 +3159,11 @@ Script should contain one command per line, just like command would be typed in
         self._postparsing_hooks = []
         self._precmd_hooks = []
         self._postcmd_hooks = []
+        self._cmdfinalization_hooks = []
 
     @classmethod
     def _validate_callable_param_count(cls, func, count):
+        """Ensure a function has the given number of parameters."""
         signature = inspect.signature(func)
         # validate that the callable has the right number of parameters
         nparam = len(signature.parameters)
@@ -3153,7 +3176,7 @@ Script should contain one command per line, just like command would be typed in
 
     @classmethod
     def _validate_prepostloop_callable(cls, func):
-        """Check parameter and return values for preloop and postloop hooks"""
+        """Check parameter and return types for preloop and postloop hooks."""
         cls._validate_callable_param_count(func, 0)
         # make sure there is no return notation
         signature = inspect.signature(func)
@@ -3174,7 +3197,7 @@ Script should contain one command per line, just like command would be typed in
 
     @classmethod
     def _validate_postparsing_callable(cls, func):
-        """Check parameter and return values for postparsing hooks"""
+        """Check parameter and return types for postparsing hooks"""
         cls._validate_callable_param_count(func, 1)
         signature = inspect.signature(func)
         _, param = list(signature.parameters.items())[0]
@@ -3194,6 +3217,7 @@ Script should contain one command per line, just like command would be typed in
 
     @classmethod
     def _validate_prepostcmd_hook(cls, func, data_type):
+        """Check parameter and return types for pre and post command hooks."""
         signature = inspect.signature(func)
         # validate that the callable has the right number of parameters
         cls._validate_callable_param_count(func, 1)
@@ -3229,6 +3253,26 @@ Script should contain one command per line, just like command would be typed in
         self._validate_prepostcmd_hook(func, plugin.PostcommandData)
         self._postcmd_hooks.append(func)
 
+    @classmethod
+    def _validate_cmdfinalization_callable(cls, func):
+        """Check parameter and return types for command finalization hooks."""
+        cls._validate_callable_param_count(func, 1)
+        signature = inspect.signature(func)
+        _, param = list(signature.parameters.items())[0]
+        if param.annotation != plugin.CommandFinalizationData:
+            raise TypeError("{} must have one parameter declared with type 'cmd2.plugin.CommandFinalizationData'".format(
+                func.__name__
+            ))
+        if signature.return_annotation != plugin.CommandFinalizationData:
+            raise TypeError("{} must declare return a return type of 'cmd2.plugin.CommandFinalizationData'".format(
+                func.__name__
+            ))
+        pass
+
+    def register_cmdfinalization_hook(self, func):
+        """Register a hook to be called after a command is completed, whether it completes successfully or not."""
+        self._validate_cmdfinalization_callable(func)
+        self._cmdfinalization_hooks.append(func)
 
 class History(list):
     """ A list of HistoryItems that knows how to respond to user requests. """
diff --git a/cmd2/plugin.py b/cmd2/plugin.py
index 377012b8..5c68dfb9 100644
--- a/cmd2/plugin.py
+++ b/cmd2/plugin.py
@@ -20,3 +20,4 @@ class PostcommandData():
 @attr.s
 class CommandFinalizationData():
     stop = attr.ib()
+    statement = attr.ib()
diff --git a/docs/hooks.rst b/docs/hooks.rst
index fd140331..0ec5d7e9 100644
--- a/docs/hooks.rst
+++ b/docs/hooks.rst
@@ -278,6 +278,8 @@ command finalization hook::
         def myhookmethod(self, stop, statement):
             return stop
 
+Command Finalization hooks must check whether the statement object is ``None``. There are certain circumstances where these hooks may be called before the statement has been parsed, so you can't always rely on having a statement.
+
 If any prior postparsing or precommand hook has requested the application to
 terminate, the value of the ``stop`` parameter passed to the first command
 finalization hook will be ``True``. Any command finalization hook can change the
diff --git a/tests/test_plugin.py b/tests/test_plugin.py
index cc60ce2a..fd8609c3 100644
--- a/tests/test_plugin.py
+++ b/tests/test_plugin.py
@@ -25,6 +25,7 @@ class Plugin:
         self.called_postparsing = 0
         self.called_precmd = 0
         self.called_postcmd = 0
+        self.called_cmdfinalization = 0
 
     ###
     #
@@ -185,6 +186,52 @@ class Plugin:
     def postcmd_hook_wrong_return_annotation(self, data: plugin.PostcommandData) -> cmd2.Statement:
         return self.statement_parser.parse('hi there')
 
+    ###
+    #
+    # command finalization hooks, some valid, some invalid
+    #
+    ###
+    def cmdfinalization_hook(self, data: plugin.CommandFinalizationData) -> plugin.CommandFinalizationData:
+        """A command finalization hook."""
+        self.called_cmdfinalization += 1
+        return data
+
+    def cmdfinalization_hook_stop(self, data: cmd2.plugin.CommandFinalizationData) -> cmd2.plugin.CommandFinalizationData:
+        "A postparsing hook which requests application exit"
+        self.called_cmdfinalization += 1
+        data.stop = True
+        return data
+
+    def cmdfinalization_hook_exception(self, data: cmd2.plugin.CommandFinalizationData) -> cmd2.plugin.CommandFinalizationData:
+        "A postparsing hook which raises an exception"
+        self.called_cmdfinalization += 1
+        raise ValueError
+
+    def cmdfinalization_hook_not_enough_parameters(self) -> plugin.CommandFinalizationData:
+        """A command finalization hook with no parameters."""
+        pass
+
+    def cmdfinalization_hook_too_many_parameters(self, one: plugin.CommandFinalizationData, two: str) -> plugin.CommandFinalizationData:
+        """A command finalization hook with too many parameters."""
+        return one
+
+    def cmdfinalization_hook_no_parameter_annotation(self, data) -> plugin.CommandFinalizationData:
+        """A command finalization hook with no type annotation on the parameter."""
+        return data
+
+    def cmdfinalization_hook_wrong_parameter_annotation(self, data: str) -> plugin.CommandFinalizationData:
+        """A command finalization hook with the incorrect type annotation on the parameter."""
+        return data
+
+    def cmdfinalization_hook_no_return_annotation(self, data: plugin.CommandFinalizationData):
+        """A command finalizationhook with no type annotation on the return value."""
+        return data
+
+    def cmdfinalization_hook_wrong_return_annotation(self, data: plugin.CommandFinalizationData) -> cmd2.Statement:
+        """A command finalization hook with the wrong return type annotation."""
+        return self.statement_parser.parse('hi there')
+
+
 class PluggedApp(Plugin, cmd2.Cmd):
     "A sample app with a plugin mixed in"
     def __init__(self, *args, **kwargs):
@@ -413,8 +460,8 @@ def test_postparsing_hook_exception(capsys):
 
     # register another function, but it shouldn't be called
     app.reset_counters()
-    stop = app.register_postparsing_hook(app.postparse_hook)
-    app.onecmd_plus_hooks('say hello')
+    app.register_postparsing_hook(app.postparse_hook)
+    stop = app.onecmd_plus_hooks('say hello')
     out, err = capsys.readouterr()
     assert not stop
     assert not out
@@ -658,6 +705,95 @@ def test_postcmd_exception_second(capsys):
 # command finalization
 #
 ###
+def test_register_cmdfinalization_hook_parameter_count():
+    app = PluggedApp()
+    with pytest.raises(TypeError):
+        app.register_cmdfinalization_hook(app.cmdfinalization_hook_not_enough_parameters)
+    with pytest.raises(TypeError):
+        app.register_cmdfinalization_hook(app.cmdfinalization_hook_too_many_parameters)
+
+def test_register_cmdfinalization_hook_no_parameter_annotation():
+    app = PluggedApp()
+    with pytest.raises(TypeError):
+        app.register_cmdfinalization_hook(app.cmdfinalization_hook_no_parameter_annotation)
+
+def test_register_cmdfinalization_hook_wrong_parameter_annotation():
+    app = PluggedApp()
+    with pytest.raises(TypeError):
+        app.register_cmdfinalization_hook(app.cmdfinalization_hook_wrong_parameter_annotation)
+
+def test_register_cmdfinalization_hook_no_return_annotation():
+    app = PluggedApp()
+    with pytest.raises(TypeError):
+        app.register_cmdfinalization_hook(app.cmdfinalization_hook_no_return_annotation)
+
+def test_register_cmdfinalization_hook_wrong_return_annotation():
+    app = PluggedApp()
+    with pytest.raises(TypeError):
+        app.register_cmdfinalization_hook(app.cmdfinalization_hook_wrong_return_annotation)
+
 def test_cmdfinalization(capsys):
-    pass
+    app = PluggedApp()
+    app.onecmd_plus_hooks('say hello')
+    out, err = capsys.readouterr()
+    assert out == 'hello\n'
+    assert not err
+    assert app.called_cmdfinalization == 0
+
+    app.register_cmdfinalization_hook(app.cmdfinalization_hook)
+    app.onecmd_plus_hooks('say hello')
+    out, err = capsys.readouterr()
+    assert out == 'hello\n'
+    assert not err
+    assert app.called_cmdfinalization == 1
 
+    # register the function again, so it should be called twice
+    app.reset_counters()
+    app.register_cmdfinalization_hook(app.cmdfinalization_hook)
+    app.onecmd_plus_hooks('say hello')
+    out, err = capsys.readouterr()
+    assert out == 'hello\n'
+    assert not err
+    assert app.called_cmdfinalization == 2
+
+def test_cmdfinalization_stop_first(capsys):
+    app = PluggedApp()
+    app.register_cmdfinalization_hook(app.cmdfinalization_hook_stop)
+    app.register_cmdfinalization_hook(app.cmdfinalization_hook)
+    stop = app.onecmd_plus_hooks('say hello')
+    out, err = capsys.readouterr()
+    assert out == 'hello\n'
+    assert not err
+    assert app.called_cmdfinalization == 2
+    assert stop
+
+def test_cmdfinalization_stop_second(capsys):
+    app = PluggedApp()
+    app.register_cmdfinalization_hook(app.cmdfinalization_hook)
+    app.register_cmdfinalization_hook(app.cmdfinalization_hook_stop)
+    stop = app.onecmd_plus_hooks('say hello')
+    out, err = capsys.readouterr()
+    assert out == 'hello\n'
+    assert not err
+    assert app.called_cmdfinalization == 2
+    assert stop
+
+def test_cmdfinalization_hook_exception(capsys):
+    app = PluggedApp()
+    app.register_cmdfinalization_hook(app.cmdfinalization_hook_exception)
+    stop = app.onecmd_plus_hooks('say hello')
+    out, err = capsys.readouterr()
+    assert not stop
+    assert out == 'hello\n'
+    assert err
+    assert app.called_cmdfinalization == 1
+
+    # register another function, but it shouldn't be called
+    app.reset_counters()
+    app.register_cmdfinalization_hook(app.cmdfinalization_hook)
+    stop = app.onecmd_plus_hooks('say hello')
+    out, err = capsys.readouterr()
+    assert not stop
+    assert out == 'hello\n'
+    assert err
+    assert app.called_cmdfinalization == 1