Postcommand hooks implemented

4 files changed, 207 insertions, 50 deletions

diff --git a/cmd2/cmd2.py b/cmd2/cmd2.py
index bd56fe71..0bb7509f 100644
--- a/cmd2/cmd2.py
+++ b/cmd2/cmd2.py
@@ -1726,8 +1726,12 @@ class Cmd(cmd.Cmd):
                 stop = self.onecmd(statement)
 
                 # postcommand hooks
+                data = plugin.PostcommandData(stop, statement)
                 for func in self._postcmd_hooks:
-                    stop = func(stop, statement)
+                    data = func(data)
+                # retrieve the final value of stop, ignoring any statement modification from the hooks
+                stop = data.stop
+                # call postcmd() for compatibility with cmd.Cmd
                 stop = self.postcmd(stop, statement)
 
                 if self.timing:
@@ -1738,6 +1742,7 @@ class Cmd(cmd.Cmd):
         except EmptyStatement:
             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:
@@ -3187,38 +3192,42 @@ Script should contain one command per line, just like command would be typed in
         self._validate_postparsing_callable(func)
         self._postparsing_hooks.append(func)
 
-    def register_precmd_hook(self, func):
-        """Register a function to be called before the command function."""
+    @classmethod
+    def _validate_prepostcmd_hook(cls, func, data_type):
         signature = inspect.signature(func)
         # validate that the callable has the right number of parameters
-        self._validate_callable_param_count(func, 1)
+        cls._validate_callable_param_count(func, 1)
         # validate the parameter has the right annotation
         paramname = list(signature.parameters.keys())[0]
         param = signature.parameters[paramname]
-        if param.annotation != plugin.PrecommandData:
+        if param.annotation != data_type:
             raise TypeError('argument 1 of {} has incompatible type {}, expected {}'.format(
                 func.__name__,
                 param.annotation,
-                plugin.PrecommandData,
+                data_type,
             ))
         # validate the return value has the right annotation
         if signature.return_annotation == signature.empty:
             raise TypeError('{} does not have a declared return type, expected {}'.format(
                 func.__name__,
-                plugin.PrecommandData,
+                data_type,
             ))
-        if signature.return_annotation != plugin.PrecommandData:
+        if signature.return_annotation != data_type:
             raise TypeError('{} has incompatible return type {}, expected {}'.format(
                 func.__name__,
                 signature.return_annotation,
-                plugin.PrecommandData,
+                data_type,
             ))
+
+    def register_precmd_hook(self, func):
+        """Register a hook to be called before the command function."""
+        self._validate_prepostcmd_hook(func, plugin.PrecommandData)
         self._precmd_hooks.append(func)
 
     def register_postcmd_hook(self, func):
-        """Register a function to be called after the command function."""
+        """Register a hook to be called after the command function."""
+        self._validate_prepostcmd_hook(func, plugin.PostcommandData)
         self._postcmd_hooks.append(func)
-        # TODO check signature of registered func and throw error if it's wrong
 
 
 class History(list):
diff --git a/cmd2/plugin.py b/cmd2/plugin.py
index 47f4a514..377012b8 100644
--- a/cmd2/plugin.py
+++ b/cmd2/plugin.py
@@ -14,8 +14,9 @@ class PrecommandData():
 
 @attr.s
 class PostcommandData():
-    stop = attr.ib(default=False)
+    stop = attr.ib()
+    statement = attr.ib()
 
 @attr.s
 class CommandFinalizationData():
-    stop = attr.ib(default=False)
+    stop = attr.ib()
diff --git a/docs/hooks.rst b/docs/hooks.rst
index 4c15e90c..fd140331 100644
--- a/docs/hooks.rst
+++ b/docs/hooks.rst
@@ -196,10 +196,10 @@ Once output is redirected and the timer started, all the hooks registered with
             super().__init__(*args, **kwargs)
             self.register_precmd_hook(self.myhookmethod)
 
-        def myhookmethod(self, params: cmd2.plugin.PrecommandData) -> cmd2.plugin.PrecommandData:
+        def myhookmethod(self, data: cmd2.plugin.PrecommandData) -> cmd2.plugin.PrecommandData:
             # the statement object created from the user input
-            # is available as params.statement
-            return params
+            # is available as data.statement
+            return data
 
 ``register_precmd_hook()`` checks the method signature of the passed callable,
 and raises a ``TypeError`` if it has the wrong number of parameters. It will
@@ -208,7 +208,7 @@ as ``PrecommandData``.
 
 You may choose to modify the user input by creating a new ``Statement`` with
 different properties (see above). If you do so, assign your new ``Statement``
-object to ``params.statement``.
+object to ``data.statement``.
 
 The precommand hook must return a ``PrecommandData`` object. You don't have to
 create this object from scratch, you can just return the one passed into the hook.
@@ -227,22 +227,23 @@ method`` has been called and returns, all postcommand hooks are called. If
 output was redirected by the user, it is still redirected, and the command timer
 is still running.
 
-Here's how to define a register a postcommand hook::
+Here's how to define and register a postcommand hook::
 
     class App(cmd2.Cmd):
         def __init__(self, *args, *kwargs):
             super().__init__(*args, **kwargs)
             self.register_postcmd_hook(self.myhookmethod)
 
-        def myhookmethod(self, stop, statement):
+        def myhookmethod(self, data: cmd2.plugin.PostcommandData) -> cmd2.plugin.PostcommandData:
             return stop
 
-Your hook will be passed the statement object, which describes the command which
-was executed. If your postcommand hook method gets called, you are guaranteed
-that the command method was called, and that it didn't raise an exception.
+Your hook will be passed a ``PostcommandData`` object, which has a ``statement``
+attribute that describes the command which was executed. If your postcommand
+hook method gets called, you are guaranteed that the command method was called,
+and that it didn't raise an exception.
 
-If any postcommand hook raises an exception, no further postcommand hook methods
-will be called.
+If any postcommand hook raises an exception, the exception will be displayed to the user, and no further postcommand hook methods
+will be called. Command finalization hooks, if any, will be called.
 
 After all registered precommand hooks have been called,
 ``self.postcmd(statement)`` will be called to retain full backward compatibility
diff --git a/tests/test_plugin.py b/tests/test_plugin.py
index 1f07f177..cc60ce2a 100644
--- a/tests/test_plugin.py
+++ b/tests/test_plugin.py
@@ -24,6 +24,7 @@ class Plugin:
     def reset_counters(self):
         self.called_postparsing = 0
         self.called_precmd = 0
+        self.called_postcmd = 0
 
     ###
     #
@@ -43,7 +44,7 @@ class Plugin:
         "A preloop or postloop hook with too many parameters"
         pass
 
-    def prepost_hook_with_wrong_return_type(self) -> bool:
+    def prepost_hook_with_wrong_return_annotation(self) -> bool:
         "A preloop or postloop hook with incorrect return type"
         pass
 
@@ -77,19 +78,19 @@ class Plugin:
         "A postparsing hook with too many parameters"
         pass
 
-    def postparse_hook_undeclared_parameter_type(self, data) -> cmd2.plugin.PostparsingData:
+    def postparse_hook_undeclared_parameter_annotation(self, data) -> cmd2.plugin.PostparsingData:
         "A postparsing hook with an undeclared parameter type"
         pass
 
-    def postparse_hook_wrong_parameter_type(self, data: str) -> cmd2.plugin.PostparsingData:
+    def postparse_hook_wrong_parameter_annotation(self, data: str) -> cmd2.plugin.PostparsingData:
         "A postparsing hook with the wrong parameter type"
         pass
 
-    def postparse_hook_undeclared_return_type(self, data: cmd2.plugin.PostparsingData):
+    def postparse_hook_undeclared_return_annotation(self, data: cmd2.plugin.PostparsingData):
         "A postparsing hook with an undeclared return type"
         pass
 
-    def postparse_hook_wrong_return_type(self, data: cmd2.plugin.PostparsingData) -> str:
+    def postparse_hook_wrong_return_annotation(self, data: cmd2.plugin.PostparsingData) -> str:
         "A postparsing hook with the wrong return type"
         pass
 
@@ -126,21 +127,63 @@ class Plugin:
         "A precommand hook with too many parameters"
         return one
 
-    def precmd_hook_no_parameter_type(self, data) -> plugin.PrecommandData:
+    def precmd_hook_no_parameter_annotation(self, data) -> plugin.PrecommandData:
         "A precommand hook with no type annotation on the parameter"
         return data
 
-    def precmd_hook_wrong_parameter_type(self, data: str) -> plugin.PrecommandData:
+    def precmd_hook_wrong_parameter_annotation(self, data: str) -> plugin.PrecommandData:
         "A precommand hook with the incorrect type annotation on the parameter"
         return data
 
-    def precmd_hook_no_return_type(self, data: plugin.PrecommandData):
+    def precmd_hook_no_return_annotation(self, data: plugin.PrecommandData):
         "A precommand hook with no type annotation on the return value"
         return data
 
-    def precmd_hook_wrong_return_type(self, data: plugin.PrecommandData) -> cmd2.Statement:
+    def precmd_hook_wrong_return_annotation(self, data: plugin.PrecommandData) -> cmd2.Statement:
         return self.statement_parser.parse('hi there')
 
+    ###
+    #
+    # postcommand hooks, some valid, some invalid
+    #
+    ###
+    def postcmd(self, stop: bool, statement: cmd2.Statement) -> bool:
+        "Override cmd.Cmd method"
+        self.called_postcmd += 1
+        return stop
+
+    def postcmd_hook(self, data: plugin.PostcommandData) -> plugin.PostcommandData:
+        "A postcommand hook"
+        self.called_postcmd += 1
+        return data
+
+    def postcmd_hook_exception(self, data: plugin.PostcommandData) -> plugin.PostcommandData:
+        "A postcommand hook with raises an exception"
+        self.called_postcmd += 1
+        raise ZeroDivisionError
+
+    def postcmd_hook_not_enough_parameters(self) -> plugin.PostcommandData:
+        "A precommand hook with no parameters"
+        pass
+
+    def postcmd_hook_too_many_parameters(self, one: plugin.PostcommandData, two: str) -> plugin.PostcommandData:
+        "A precommand hook with too many parameters"
+        return one
+
+    def postcmd_hook_no_parameter_annotation(self, data) -> plugin.PostcommandData:
+        "A precommand hook with no type annotation on the parameter"
+        return data
+
+    def postcmd_hook_wrong_parameter_annotation(self, data: str) -> plugin.PostcommandData:
+        "A precommand hook with the incorrect type annotation on the parameter"
+        return data
+
+    def postcmd_hook_no_return_annotation(self, data: plugin.PostcommandData):
+        "A precommand hook with no type annotation on the return value"
+        return data
+
+    def postcmd_hook_wrong_return_annotation(self, data: plugin.PostcommandData) -> cmd2.Statement:
+        return self.statement_parser.parse('hi there')
 
 class PluggedApp(Plugin, cmd2.Cmd):
     "A sample app with a plugin mixed in"
@@ -161,10 +204,10 @@ def test_register_preloop_hook_too_many_parameters():
     with pytest.raises(TypeError):
         app.register_preloop_hook(app.prepost_hook_too_many_parameters)
 
-def test_register_preloop_hook_with_return_type():
+def test_register_preloop_hook_with_return_annotation():
     app = PluggedApp()
     with pytest.raises(TypeError):
-        app.register_preloop_hook(app.prepost_hook_with_wrong_return_type)
+        app.register_preloop_hook(app.prepost_hook_with_wrong_return_annotation)
 
 def test_preloop_hook(capsys):
     app = PluggedApp()
@@ -192,10 +235,10 @@ def test_register_postloop_hook_too_many_parameters():
     with pytest.raises(TypeError):
         app.register_postloop_hook(app.prepost_hook_too_many_parameters)
 
-def test_register_postloop_hook_with_wrong_return_type():
+def test_register_postloop_hook_with_wrong_return_annotation():
     app = PluggedApp()
     with pytest.raises(TypeError):
-        app.register_postloop_hook(app.prepost_hook_with_wrong_return_type)
+        app.register_postloop_hook(app.prepost_hook_with_wrong_return_annotation)
 
 def test_postloop_hook(capsys):
     app = PluggedApp()
@@ -228,25 +271,25 @@ def test_postparsing_hook_too_many_parameters():
     with pytest.raises(TypeError):
         app.register_postparsing_hook(app.postparse_hook_too_many_parameters)
 
-def test_postparsing_hook_undeclared_parameter_type():
+def test_postparsing_hook_undeclared_parameter_annotation():
     app = PluggedApp()
     with pytest.raises(TypeError):
-        app.register_postparsing_hook(app.postparse_hook_undeclared_parameter_type)
+        app.register_postparsing_hook(app.postparse_hook_undeclared_parameter_annotation)
 
-def test_postparsing_hook_wrong_parameter_type():
+def test_postparsing_hook_wrong_parameter_annotation():
     app = PluggedApp()
     with pytest.raises(TypeError):
-        app.register_postparsing_hook(app.postparse_hook_wrong_parameter_type)
+        app.register_postparsing_hook(app.postparse_hook_wrong_parameter_annotation)
 
-def test_postparsing_hook_undeclared_return_type():
+def test_postparsing_hook_undeclared_return_annotation():
     app = PluggedApp()
     with pytest.raises(TypeError):
-        app.register_postparsing_hook(app.postparse_hook_undeclared_return_type)
+        app.register_postparsing_hook(app.postparse_hook_undeclared_return_annotation)
 
-def test_postparsing_hook_wrong_return_type():
+def test_postparsing_hook_wrong_return_annotation():
     app = PluggedApp()
     with pytest.raises(TypeError):
-        app.register_postparsing_hook(app.postparse_hook_wrong_return_type)
+        app.register_postparsing_hook(app.postparse_hook_wrong_return_annotation)
 
 def test_postparsing_hook(capsys):
     app = PluggedApp()
@@ -393,22 +436,22 @@ def test_register_precmd_hook_parameter_count():
 def test_register_precmd_hook_no_parameter_annotation():
     app = PluggedApp()
     with pytest.raises(TypeError):
-        app.register_precmd_hook(app.precmd_hook_no_parameter_type)
+        app.register_precmd_hook(app.precmd_hook_no_parameter_annotation)
 
 def test_register_precmd_hook_wrong_parameter_annotation():
     app = PluggedApp()
     with pytest.raises(TypeError):
-        app.register_precmd_hook(app.precmd_hook_wrong_parameter_type)
+        app.register_precmd_hook(app.precmd_hook_wrong_parameter_annotation)
 
 def test_register_precmd_hook_no_return_annotation():
     app = PluggedApp()
     with pytest.raises(TypeError):
-        app.register_precmd_hook(app.precmd_hook_no_return_type)
+        app.register_precmd_hook(app.precmd_hook_no_return_annotation)
 
 def test_register_precmd_hook_wrong_return_annotation():
     app = PluggedApp()
     with pytest.raises(TypeError):
-        app.register_precmd_hook(app.precmd_hook_wrong_return_type)
+        app.register_precmd_hook(app.precmd_hook_wrong_return_annotation)
 
 def test_precmd_hook(capsys):
     app = PluggedApp()
@@ -504,8 +547,111 @@ def test_precmd_hook_emptystatement_second(capsys):
 # test postcmd hooks
 #
 ####
+def test_register_postcmd_hook_parameter_count():
+    app = PluggedApp()
+    with pytest.raises(TypeError):
+        app.register_postcmd_hook(app.postcmd_hook_not_enough_parameters)
+    with pytest.raises(TypeError):
+        app.register_postcmd_hook(app.postcmd_hook_too_many_parameters)
+
+def test_register_postcmd_hook_no_parameter_annotation():
+    app = PluggedApp()
+    with pytest.raises(TypeError):
+        app.register_postcmd_hook(app.postcmd_hook_no_parameter_annotation)
+
+def test_register_postcmd_hook_wrong_parameter_annotation():
+    app = PluggedApp()
+    with pytest.raises(TypeError):
+        app.register_postcmd_hook(app.postcmd_hook_wrong_parameter_annotation)
+
+def test_register_postcmd_hook_no_return_annotation():
+    app = PluggedApp()
+    with pytest.raises(TypeError):
+        app.register_postcmd_hook(app.postcmd_hook_no_return_annotation)
+
+def test_register_postcmd_hook_wrong_return_annotation():
+    app = PluggedApp()
+    with pytest.raises(TypeError):
+        app.register_postcmd_hook(app.postcmd_hook_wrong_return_annotation)
+
 def test_postcmd(capsys):
-    pass
+    app = PluggedApp()
+    app.onecmd_plus_hooks('say hello')
+    out, err = capsys.readouterr()
+    assert out == 'hello\n'
+    assert not err
+    # without registering any hooks, postcmd() should be called
+    assert app.called_postcmd == 1
+
+    app.reset_counters()
+    app.register_postcmd_hook(app.postcmd_hook)
+    app.onecmd_plus_hooks('say hello')
+    out, err = capsys.readouterr()
+    assert out == 'hello\n'
+    assert not err
+    # with one hook registered, we should get precmd() and the hook
+    assert app.called_postcmd == 2
+
+    # register the function again, so it should be called twice
+    app.reset_counters()
+    app.register_postcmd_hook(app.postcmd_hook)
+    app.onecmd_plus_hooks('say hello')
+    out, err = capsys.readouterr()
+    assert out == 'hello\n'
+    assert not err
+    # with two hooks registered, we should get precmd() and both hooks
+    assert app.called_postcmd == 3
+
+def test_postcmd_exception_first(capsys):
+    app = PluggedApp()
+    app.register_postcmd_hook(app.postcmd_hook_exception)
+    stop = app.onecmd_plus_hooks('say hello')
+    out, err = capsys.readouterr()
+    assert not stop
+    assert out == 'hello\n'
+    assert err
+    # since the registered hooks are called before postcmd(), if a registered
+    # hook throws an exception, postcmd() is never called. So we should have
+    # a count of one because we called the hook that raised the exception
+    assert app.called_postcmd == 1
+
+    # register another function but it shouldn't be called
+    app.reset_counters()
+    stop = app.register_postcmd_hook(app.postcmd_hook)
+    app.onecmd_plus_hooks('say hello')
+    out, err = capsys.readouterr()
+    assert not stop
+    assert out == 'hello\n'
+    assert err
+    # the exception raised by the first hook should prevent the second
+    # hook from being called, and it also prevents postcmd() from being
+    # called
+    assert app.called_postcmd == 1
+
+def test_postcmd_exception_second(capsys):
+    app = PluggedApp()
+    app.register_postcmd_hook(app.postcmd_hook)
+    stop = app.onecmd_plus_hooks('say hello')
+    out, err = capsys.readouterr()
+    assert not stop
+    assert out == 'hello\n'
+    assert not err
+    # with one hook registered, we should get the hook and postcmd()
+    assert app.called_postcmd == 2
+
+    # register another function which should be called
+    app.reset_counters()
+    stop = app.register_postcmd_hook(app.postcmd_hook_exception)
+    app.onecmd_plus_hooks('say hello')
+    out, err = capsys.readouterr()
+    assert not stop
+    assert out == 'hello\n'
+    assert err
+    # the exception raised by the first hook should prevent the second
+    # hook from being called, and it also prevents postcmd() from being
+    # called. So we have the first hook, and the second hook, which raised
+    # the exception
+    assert app.called_postcmd == 2
 
 ##
 #