File: test_output_types.py

package info (click to toggle)
python-nbstripout 0.9.0-1~exp2
  • links: PTS, VCS
  • area: main
  • in suites: experimental
  • size: 464 kB
  • sloc: python: 1,068; sh: 20; makefile: 13
file content (92 lines) | stat: -rw-r--r-- 3,012 bytes parent folder | download | duplicates (2)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
import os
from copy import deepcopy

import nbformat
import pytest

from nbstripout import strip_output

directory = os.path.dirname(__file__)


@pytest.fixture
def orig_nb():
    # Grab the original notebook
    fname = 'test_output_types.ipynb'
    return nbformat.read(os.path.join(directory, fname), nbformat.NO_CONVERT)


def test_drop_errors(orig_nb):
    """
    Confirm that --drop-output-types works as expected,
    when asking just to drop `error` outputs.
    """
    nb_stripped = strip_output(
        deepcopy(orig_nb), keep_output=True, keep_count=False, keep_id=False, drop_output_types={'error'}
    )

    # No outputs in the markdown
    assert not hasattr(nb_stripped.cells[0], 'outputs')

    # Original cell should have 3 outputs, with the last being error
    assert len(orig_nb.cells[1].outputs) == 3
    assert orig_nb.cells[1].outputs[2]['output_type'] == 'error'

    # Second cell should have a stdout stream, stderr stream and an error
    stripped_output_1 = nb_stripped.cells[1].outputs[0]
    stripped_output_2 = nb_stripped.cells[1].outputs[1]
    assert len(nb_stripped.cells[1].outputs) == 2
    assert stripped_output_1['output_type'] == 'stream'
    assert stripped_output_2['output_type'] == 'stream'
    assert stripped_output_1['name'] == 'stdout'
    assert stripped_output_2['name'] == 'stderr'

    # Third cell should have an execution output
    assert len(nb_stripped.cells[2].outputs) == 1
    assert nb_stripped.cells[2].outputs[0]['output_type'] == 'execute_result'


def test_keep_output(orig_nb):
    """
    Confirm that --keep-output-types works as expected,
    dropping all but `execute_result` outputs from the notebook.
    """
    nb_stripped = strip_output(
        deepcopy(orig_nb), keep_output=True, keep_count=False, keep_id=False, keep_output_types={'execute_result'}
    )

    # No outputs in the markdown
    assert not hasattr(nb_stripped.cells[0], 'outputs')

    # Original cell should have 3 outputs, with the last being error
    assert len(orig_nb.cells[1].outputs) == 3
    assert orig_nb.cells[1].outputs[2]['output_type'] == 'error'

    # All outputs should be stripped in the second cell
    assert len(nb_stripped.cells[1].outputs) == 0

    # Third cell should have an execution output
    assert len(nb_stripped.cells[2].outputs) == 1


def test_output_format_tags(orig_nb):
    """
    Confirm that both <output_type>:<name> and <output_type> formats work.
    """
    nb_stripped = strip_output(
        deepcopy(orig_nb),
        keep_output=False,
        keep_count=False,
        keep_id=False,
        keep_output_types={'stream:stdout', 'execute_result'},
    )

    # No outputs in the markdown
    assert not hasattr(nb_stripped.cells[0], 'outputs')

    # Stripping all but stdout should leave only the print statement
    assert len(orig_nb.cells[1].outputs) == 3
    assert len(nb_stripped.cells[1].outputs) == 1

    # Third cell should have only the execute_result
    assert len(nb_stripped.cells[2].outputs) == 1