Extend unit tests for the Output widget

ipywidgets/widgets/tests/test_widget_output.py

@@ -1,50 +1,193 @@
 import sys
+from unittest import TestCase
+from contextlib import contextmanager
 
 from IPython.display import Markdown, Image
-from ipywidgets import Output
+from ipywidgets import widget_output
+
+
+class TestOutputWidget(TestCase):
+
+    @contextmanager
+    def _mocked_ipython(self, get_ipython, clear_output):
+        """ Context manager that monkeypatches get_ipython and clear_output """
+        original_clear_output = widget_output.clear_output
+        original_get_ipython = widget_output.get_ipython
+        widget_output.get_ipython = get_ipython
+        widget_output.clear_output = clear_output
+        try:
+            yield
+        finally:
+            widget_output.clear_output = original_clear_output
+            widget_output.get_ipython = original_get_ipython
+
+    def _mock_get_ipython(self, msg_id):
+        """ Returns a mock IPython application with a mocked kernel """
+        kernel = type(
+            'mock_kernel',
+            (object, ),
+            {'_parent_header': {'header': {'msg_id': msg_id}}}
+        )
+
+        # Specifically override this so the traceback
+        # is still printed to screen
+        def showtraceback(self_, exc_tuple, *args, **kwargs):
+            etype, evalue, tb = exc_tuple
+            raise etype(evalue)
+
+        ipython = type(
+            'mock_ipython',
+            (object, ),
+            {'kernel': kernel, 'showtraceback': showtraceback}
+        )
+        return ipython
+
+    def _mock_clear_output(self):
+        """ Mock function that records calls to it """
+        calls = []
+
+        def clear_output(*args, **kwargs):
+            calls.append((args, kwargs))
+        clear_output.calls = calls
+
+        return clear_output
+
+    def test_set_msg_id_when_capturing(self):
+        msg_id = 'msg-id'
+        get_ipython = self._mock_get_ipython(msg_id)
+        clear_output = self._mock_clear_output()
+
+        with self._mocked_ipython(get_ipython, clear_output):
+            widget = widget_output.Output()
+            assert widget.msg_id == ''
+            with widget:
+                assert widget.msg_id == msg_id
+            assert widget.msg_id == ''
+
+    def test_clear_output(self):
+        msg_id = 'msg-id'
+        get_ipython = self._mock_get_ipython(msg_id)
+        clear_output = self._mock_clear_output()
+
+        with self._mocked_ipython(get_ipython, clear_output):
+            widget = widget_output.Output()
+            widget.clear_output(wait=True)
+
+        assert len(clear_output.calls) == 1
+        assert clear_output.calls[0] == ((), {'wait': True})
+
+    def test_capture_decorator(self):
+        msg_id = 'msg-id'
+        get_ipython = self._mock_get_ipython(msg_id)
+        clear_output = self._mock_clear_output()
+        expected_argument = 'arg'
+        expected_keyword_argument = True
+        captee_calls = []
+
+        with self._mocked_ipython(get_ipython, clear_output):
+            widget = widget_output.Output()
+            assert widget.msg_id == ''
+
+            @widget.capture()
+            def captee(*args, **kwargs):
+                # Check that we are capturing output
+                assert widget.msg_id == msg_id
+
+                # Check that arguments are passed correctly
+                captee_calls.append((args, kwargs))
+
+            captee(
+                expected_argument, keyword_argument=expected_keyword_argument)
+            assert widget.msg_id == ''
+            captee()
+
+        assert len(captee_calls) == 2
+        assert captee_calls[0] == (
+            (expected_argument, ),
+            {'keyword_argument': expected_keyword_argument}
+        )
+        assert captee_calls[1] == ((), {})
+
+    def test_capture_decorator_clear_output(self):
+        msg_id = 'msg-id'
+        get_ipython = self._mock_get_ipython(msg_id)
+        clear_output = self._mock_clear_output()
+
+        with self._mocked_ipython(get_ipython, clear_output):
+            widget = widget_output.Output()
+
+            @widget.capture(clear_output=True, wait=True)
+            def captee(*args, **kwargs):
+                # Check that we are capturing output
+                assert widget.msg_id == msg_id
+
+            captee()
+            captee()
+
+        assert len(clear_output.calls) == 2
+        assert clear_output.calls[0] == clear_output.calls[1] == \
+            ((), {'wait': True})
+
+    def test_capture_decorator_no_clear_output(self):
+        msg_id = 'msg-id'
+        get_ipython = self._mock_get_ipython(msg_id)
+        clear_output = self._mock_clear_output()
+
+        with self._mocked_ipython(get_ipython, clear_output):
+            widget = widget_output.Output()
+
+            @widget.capture(clear_output=False)
+            def captee(*args, **kwargs):
+                # Check that we are capturing output
+                assert widget.msg_id == msg_id
+
+            captee()
+            captee()
+
+        assert len(clear_output.calls) == 0
 
 
 def _make_stream_output(text, name):
-        return {
-            'output_type': 'stream',
-            'name': name,
-            'text': text
-        }
+    return {
+        'output_type': 'stream',
+        'name': name,
+        'text': text
+    }
 
 
 def test_append_stdout():
-    output = Output()
+    widget = widget_output.Output()
 
     # Try appending a message to stdout.
-    output.append_stdout("snakes!")
+    widget.append_stdout("snakes!")
     expected = (_make_stream_output("snakes!", "stdout"),)
-    assert output.outputs == expected, repr(output.outputs)
+    assert widget.outputs == expected, repr(widget.outputs)
 
     # Try appending a second message.
-    output.append_stdout("more snakes!")
+    widget.append_stdout("more snakes!")
     expected += (_make_stream_output("more snakes!", "stdout"),)
-    assert output.outputs == expected, repr(output.outputs)
+    assert widget.outputs == expected, repr(widget.outputs)
 
 
 def test_append_stderr():
-    output = Output()
+    widget = widget_output.Output()
 
     # Try appending a message to stderr.
-    output.append_stderr("snakes!")
+    widget.append_stderr("snakes!")
     expected = (_make_stream_output("snakes!", "stderr"),)
-    assert output.outputs == expected, repr(output.outputs)
+    assert widget.outputs == expected, repr(widget.outputs)
 
     # Try appending a second message.
-    output.append_stderr("more snakes!")
+    widget.append_stderr("more snakes!")
     expected += (_make_stream_output("more snakes!", "stderr"),)
-    assert output.outputs == expected, repr(output.outputs)
+    assert widget.outputs == expected, repr(widget.outputs)
 
 
 def test_append_display_data():
-    output = Output()
+    widget = widget_output.Output()
 
     # Try appending a Markdown object.
-    output.append_display_data(Markdown("# snakes!"))
+    widget.append_display_data(Markdown("# snakes!"))
     expected = (
         {
             'output_type': 'display_data',
@@ -55,13 +198,13 @@ def test_append_display_data():
             'metadata': {}
         },
     )
-    assert output.outputs == expected, repr(output.outputs)
+    assert widget.outputs == expected, repr(widget.outputs)
 
     # Now try appending an Image.
     image_data = b"foobar"
     image_data_b64 = image_data if sys.version_info[0] < 3 else 'Zm9vYmFy\n'
 
-    output.append_display_data(Image(image_data, width=123, height=456))
+    widget.append_display_data(Image(image_data, width=123, height=456))
     expected += (
         {
             'output_type': 'display_data',
@@ -77,4 +220,4 @@ def test_append_display_data():
             }
         },
     )
-    assert output.outputs == expected, repr(output.outputs)
+    assert widget.outputs == expected, repr(widget.outputs)

ipywidgets/widgets/widget_output.py

@@ -6,12 +6,13 @@
 Represents a widget that can be used to display output within the widget area.
 """
 
+import sys
+from functools import wraps
+
 from .domwidget import DOMWidget
 from .widget import register
-from .widget_core import CoreWidget
 from .._version import __jupyter_widgets_output_version__
 
-import sys
 from traitlets import Unicode, Tuple
 from IPython.core.interactiveshell import InteractiveShell
 from IPython.display import clear_output
@@ -23,9 +24,15 @@ class Output(DOMWidget):
     """Widget used as a context manager to display output.
 
     This widget can capture and display stdout, stderr, and rich output.  To use
-    it, create an instance of it and display it.  Then use it as a context
-    manager.  Any output produced while in it's context will be captured and
-    displayed in it instead of the standard output area.
+    it, create an instance of it and display it.
+
+    You can then use it as a context manager: any output produced while in it's
+    context will be captured and displayed in it instead of the standard output
+    area.
+
+    You can also use it to decorate a function or a method. Any output produced
+    by the function will then go to the output widget. This is useful for
+    debugging widget callbacks, for instance.
 
     Example::
         import ipywidgets as widgets
@@ -37,6 +44,10 @@ class Output(DOMWidget):
 
         with out:
             print('prints to output widget')
+
+        @out.capture()
+        def func():
+            print('prints to output widget')
     """
     _view_name = Unicode('OutputView').tag(sync=True)
     _model_name = Unicode('OutputModel').tag(sync=True)
@@ -49,9 +60,47 @@ class Output(DOMWidget):
     outputs = Tuple(help="The output messages synced from the frontend.").tag(sync=True)
 
     def clear_output(self, *pargs, **kwargs):
+        """
+        Clear the content of the output widget.
+
+        Parameters
+        ----------
+
+        wait: bool
+            If True, wait to clear the output until new output is
+            available to replace it. Default: False
+        """
         with self:
             clear_output(*pargs, **kwargs)
 
+    # PY3: Force passing clear_output and clear_kwargs as kwargs
+    def capture(self, clear_output=False, *clear_args, **clear_kwargs):
+        """
+        Decorator to capture the stdout and stderr of a function.
+
+        Parameters
+        ----------
+
+        clear_output: bool
+            If True, clear the content of the output widget at every
+            new function call. Default: False
+
+        wait: bool
+            If True, wait to clear the output until new output is
+            available to replace it. This is only used if clear_output
+            is also True.
+            Default: False
+        """
+        def capture_decorator(func):
+            @wraps(func)
+            def inner(*args, **kwargs):
+                if clear_output:
+                    self.clear_output(*clear_args, **clear_kwargs)
+                with self:
+                    return func(*args, **kwargs)
+            return inner
+        return capture_decorator
+
     def __enter__(self):
         """Called upon entering output widget context manager."""
         self._flush()