1import sys
2from unittest import TestCase
3from contextlib import contextmanager
4
5from IPython.display import Markdown, Image
6from ipywidgets import widget_output
7
8
9class TestOutputWidget(TestCase):
10
11    @contextmanager
12    def _mocked_ipython(self, get_ipython, clear_output):
13        """ Context manager that monkeypatches get_ipython and clear_output """
14        original_clear_output = widget_output.clear_output
15        original_get_ipython = widget_output.get_ipython
16        widget_output.get_ipython = get_ipython
17        widget_output.clear_output = clear_output
18        try:
19            yield
20        finally:
21            widget_output.clear_output = original_clear_output
22            widget_output.get_ipython = original_get_ipython
23
24    def _mock_get_ipython(self, msg_id):
25        """ Returns a mock IPython application with a mocked kernel """
26        kernel = type(
27            'mock_kernel',
28            (object, ),
29            {'_parent_header': {'header': {'msg_id': msg_id}}}
30        )
31
32        # Specifically override this so the traceback
33        # is still printed to screen
34        def showtraceback(self_, exc_tuple, *args, **kwargs):
35            etype, evalue, tb = exc_tuple
36            raise etype(evalue)
37
38        ipython = type(
39            'mock_ipython',
40            (object, ),
41            {'kernel': kernel, 'showtraceback': showtraceback}
42        )
43        return ipython
44
45    def _mock_clear_output(self):
46        """ Mock function that records calls to it """
47        calls = []
48
49        def clear_output(*args, **kwargs):
50            calls.append((args, kwargs))
51        clear_output.calls = calls
52
53        return clear_output
54
55    def test_set_msg_id_when_capturing(self):
56        msg_id = 'msg-id'
57        get_ipython = self._mock_get_ipython(msg_id)
58        clear_output = self._mock_clear_output()
59
60        with self._mocked_ipython(get_ipython, clear_output):
61            widget = widget_output.Output()
62            assert widget.msg_id == ''
63            with widget:
64                assert widget.msg_id == msg_id
65            assert widget.msg_id == ''
66
67    def test_clear_output(self):
68        msg_id = 'msg-id'
69        get_ipython = self._mock_get_ipython(msg_id)
70        clear_output = self._mock_clear_output()
71
72        with self._mocked_ipython(get_ipython, clear_output):
73            widget = widget_output.Output()
74            widget.clear_output(wait=True)
75
76        assert len(clear_output.calls) == 1
77        assert clear_output.calls[0] == ((), {'wait': True})
78
79    def test_capture_decorator(self):
80        msg_id = 'msg-id'
81        get_ipython = self._mock_get_ipython(msg_id)
82        clear_output = self._mock_clear_output()
83        expected_argument = 'arg'
84        expected_keyword_argument = True
85        captee_calls = []
86
87        with self._mocked_ipython(get_ipython, clear_output):
88            widget = widget_output.Output()
89            assert widget.msg_id == ''
90
91            @widget.capture()
92            def captee(*args, **kwargs):
93                # Check that we are capturing output
94                assert widget.msg_id == msg_id
95
96                # Check that arguments are passed correctly
97                captee_calls.append((args, kwargs))
98
99            captee(
100                expected_argument, keyword_argument=expected_keyword_argument)
101            assert widget.msg_id == ''
102            captee()
103
104        assert len(captee_calls) == 2
105        assert captee_calls[0] == (
106            (expected_argument, ),
107            {'keyword_argument': expected_keyword_argument}
108        )
109        assert captee_calls[1] == ((), {})
110
111    def test_capture_decorator_clear_output(self):
112        msg_id = 'msg-id'
113        get_ipython = self._mock_get_ipython(msg_id)
114        clear_output = self._mock_clear_output()
115
116        with self._mocked_ipython(get_ipython, clear_output):
117            widget = widget_output.Output()
118
119            @widget.capture(clear_output=True, wait=True)
120            def captee(*args, **kwargs):
121                # Check that we are capturing output
122                assert widget.msg_id == msg_id
123
124            captee()
125            captee()
126
127        assert len(clear_output.calls) == 2
128        assert clear_output.calls[0] == clear_output.calls[1] == \
129            ((), {'wait': True})
130
131    def test_capture_decorator_no_clear_output(self):
132        msg_id = 'msg-id'
133        get_ipython = self._mock_get_ipython(msg_id)
134        clear_output = self._mock_clear_output()
135
136        with self._mocked_ipython(get_ipython, clear_output):
137            widget = widget_output.Output()
138
139            @widget.capture(clear_output=False)
140            def captee(*args, **kwargs):
141                # Check that we are capturing output
142                assert widget.msg_id == msg_id
143
144            captee()
145            captee()
146
147        assert len(clear_output.calls) == 0
148
149
150def _make_stream_output(text, name):
151    return {
152        'output_type': 'stream',
153        'name': name,
154        'text': text
155    }
156
157
158def test_append_stdout():
159    widget = widget_output.Output()
160
161    # Try appending a message to stdout.
162    widget.append_stdout("snakes!")
163    expected = (_make_stream_output("snakes!", "stdout"),)
164    assert widget.outputs == expected, repr(widget.outputs)
165
166    # Try appending a second message.
167    widget.append_stdout("more snakes!")
168    expected += (_make_stream_output("more snakes!", "stdout"),)
169    assert widget.outputs == expected, repr(widget.outputs)
170
171
172def test_append_stderr():
173    widget = widget_output.Output()
174
175    # Try appending a message to stderr.
176    widget.append_stderr("snakes!")
177    expected = (_make_stream_output("snakes!", "stderr"),)
178    assert widget.outputs == expected, repr(widget.outputs)
179
180    # Try appending a second message.
181    widget.append_stderr("more snakes!")
182    expected += (_make_stream_output("more snakes!", "stderr"),)
183    assert widget.outputs == expected, repr(widget.outputs)
184
185
186def test_append_display_data():
187    widget = widget_output.Output()
188
189    # Try appending a Markdown object.
190    widget.append_display_data(Markdown("# snakes!"))
191    expected = (
192        {
193            'output_type': 'display_data',
194            'data': {
195                'text/plain': '<IPython.core.display.Markdown object>',
196                'text/markdown': '# snakes!'
197            },
198            'metadata': {}
199        },
200    )
201    assert widget.outputs == expected, repr(widget.outputs)
202
203    # Now try appending an Image.
204    image_data = b"foobar"
205    image_data_b64 = image_data if sys.version_info[0] < 3 else 'Zm9vYmFy\n'
206
207    widget.append_display_data(Image(image_data, width=123, height=456))
208    expected += (
209        {
210            'output_type': 'display_data',
211            'data': {
212                'image/png': image_data_b64,
213                'text/plain': '<IPython.core.display.Image object>'
214            },
215            'metadata': {
216                'image/png': {
217                    'width': 123,
218                    'height': 456
219                }
220            }
221        },
222    )
223    assert widget.outputs == expected, repr(widget.outputs)
224