Back to Repositories
Testing Jupyter Notebook Formatting Implementation in Black

This test suite validates Jupyter notebook formatting functionality in Black, focusing on handling IPython magics, cell formatting, and notebook structure. It ensures proper code style enforcement while preserving notebook-specific syntax and magic commands.
Test Coverage Overview

The test suite provides comprehensive coverage of Jupyter notebook formatting scenarios.

Key areas tested include:
Cell magic command preservation and formatting
IPython magic syntax handling
Notebook metadata processing
Empty cell and invalid content handling
Integration with Black’s core formatting logic
Implementation Analysis

The testing approach uses pytest fixtures and parametrization to validate multiple formatting scenarios.

Key patterns include:
Mocking IPython dependencies for isolation
Testing both fast and slow formatting modes
Validating notebook JSON structure integrity
Testing CLI integration points
Technical Details

Testing tools and configuration:
pytest for test framework
CliRunner for command-line testing
Custom fixtures for notebook content
MonkeyPatch for dependency mocking
Parameterized test cases for edge cases
Best Practices Demonstrated

The test suite exemplifies high-quality testing practices including:
Comprehensive edge case coverage
Clear test case organization
Effective use of pytest features
Proper isolation of dependencies
Validation of both success and error paths
psf/black

tests/test_ipynb.py
            
import contextlib
import pathlib
import re
from contextlib import AbstractContextManager
from contextlib import ExitStack as does_not_raise
from dataclasses import replace

import pytest
from _pytest.monkeypatch import MonkeyPatch
from click.testing import CliRunner

from black import (
    Mode,
    NothingChanged,
    format_cell,
    format_file_contents,
    format_file_in_place,
    main,
)
from black.handle_ipynb_magics import jupyter_dependencies_are_installed
from tests.util import DATA_DIR, get_case_path, read_jupyter_notebook

with contextlib.suppress(ModuleNotFoundError):
    import IPython
pytestmark = pytest.mark.jupyter
pytest.importorskip("IPython", reason="IPython is an optional dependency")
pytest.importorskip("tokenize_rt", reason="tokenize-rt is an optional dependency")

JUPYTER_MODE = Mode(is_ipynb=True)

EMPTY_CONFIG = DATA_DIR / "empty_pyproject.toml"

runner = CliRunner()


def test_noop() -> None:
    src = 'foo = "a"'
    with pytest.raises(NothingChanged):
        format_cell(src, fast=True, mode=JUPYTER_MODE)


@pytest.mark.parametrize("fast", [True, False])
def test_trailing_semicolon(fast: bool) -> None:
    src = 'foo = "a" ;'
    result = format_cell(src, fast=fast, mode=JUPYTER_MODE)
    expected = 'foo = "a";'
    assert result == expected


def test_trailing_semicolon_with_comment() -> None:
    src = 'foo = "a" ;  # bar'
    result = format_cell(src, fast=True, mode=JUPYTER_MODE)
    expected = 'foo = "a";  # bar'
    assert result == expected


def test_trailing_semicolon_with_comment_on_next_line() -> None:
    src = "import black;

# this is a comment"
    with pytest.raises(NothingChanged):
        format_cell(src, fast=True, mode=JUPYTER_MODE)


def test_trailing_semicolon_indented() -> None:
    src = "with foo:
    plot_bar();"
    with pytest.raises(NothingChanged):
        format_cell(src, fast=True, mode=JUPYTER_MODE)


def test_trailing_semicolon_noop() -> None:
    src = 'foo = "a";'
    with pytest.raises(NothingChanged):
        format_cell(src, fast=True, mode=JUPYTER_MODE)


@pytest.mark.parametrize(
    "mode",
    [
        pytest.param(JUPYTER_MODE, id="default mode"),
        pytest.param(
            replace(JUPYTER_MODE, python_cell_magics={"cust1", "cust2"}),
            id="custom cell magics mode",
        ),
    ],
)
def test_cell_magic(mode: Mode) -> None:
    src = "%%time
foo =bar"
    result = format_cell(src, fast=True, mode=mode)
    expected = "%%time
foo = bar"
    assert result == expected


def test_cell_magic_noop() -> None:
    src = "%%time
2 + 2"
    with pytest.raises(NothingChanged):
        format_cell(src, fast=True, mode=JUPYTER_MODE)


@pytest.mark.parametrize(
    "mode",
    [
        pytest.param(JUPYTER_MODE, id="default mode"),
        pytest.param(
            replace(JUPYTER_MODE, python_cell_magics={"cust1", "cust2"}),
            id="custom cell magics mode",
        ),
    ],
)
@pytest.mark.parametrize(
    "src, expected",
    (
        pytest.param("ls =!ls", "ls = !ls", id="System assignment"),
        pytest.param("!ls
'foo'", '!ls
"foo"', id="System call"),
        pytest.param("!!ls
'foo'", '!!ls
"foo"', id="Other system call"),
        pytest.param("?str
'foo'", '?str
"foo"', id="Help"),
        pytest.param("??str
'foo'", '??str
"foo"', id="Other help"),
        pytest.param(
            "%matplotlib inline
'foo'",
            '%matplotlib inline
"foo"',
            id="Line magic with argument",
        ),
        pytest.param("%time
'foo'", '%time
"foo"', id="Line magic without argument"),
        pytest.param(
            "env =  %env var", "env = %env var", id="Assignment to environment variable"
        ),
        pytest.param("env =  %env", "env = %env", id="Assignment to magic"),
    ),
)
def test_magic(src: str, expected: str, mode: Mode) -> None:
    result = format_cell(src, fast=True, mode=mode)
    assert result == expected


@pytest.mark.parametrize(
    "src",
    (
        "%%bash
2+2",
        "%%html --isolated
2+2",
        "%%writefile e.txt
  meh
 meh",
    ),
)
def test_non_python_magics(src: str) -> None:
    with pytest.raises(NothingChanged):
        format_cell(src, fast=True, mode=JUPYTER_MODE)


@pytest.mark.skipif(
    IPython.version_info < (8, 3),
    reason="Change in how TransformerManager transforms this input",
)
def test_set_input() -> None:
    src = "a = b??"
    expected = "??b"
    result = format_cell(src, fast=True, mode=JUPYTER_MODE)
    assert result == expected


def test_input_already_contains_transformed_magic() -> None:
    src = '%time foo()
get_ipython().run_cell_magic("time", "", "foo()\
")'
    with pytest.raises(NothingChanged):
        format_cell(src, fast=True, mode=JUPYTER_MODE)


def test_magic_noop() -> None:
    src = "ls = !ls"
    with pytest.raises(NothingChanged):
        format_cell(src, fast=True, mode=JUPYTER_MODE)


def test_cell_magic_with_magic() -> None:
    src = "%%timeit -n1
ls =!ls"
    result = format_cell(src, fast=True, mode=JUPYTER_MODE)
    expected = "%%timeit -n1
ls = !ls"
    assert result == expected


@pytest.mark.parametrize(
    "src, expected",
    (
        ("


%time 

", "%time"),
        ("  
\t
%%timeit -n4 \t 
x=2  
\r
", "%%timeit -n4
x = 2"),
        (
            "  \t

%%capture 
x=2 
%config 

%env
\t  
 

",
            "%%capture
x = 2
%config

%env",
        ),
    ),
)
def test_cell_magic_with_empty_lines(src: str, expected: str) -> None:
    result = format_cell(src, fast=True, mode=JUPYTER_MODE)
    assert result == expected


@pytest.mark.parametrize(
    "mode, expected_output, expectation",
    [
        pytest.param(
            JUPYTER_MODE,
            "%%custom_python_magic -n1 -n2
x=2",
            pytest.raises(NothingChanged),
            id="No change when cell magic not registered",
        ),
        pytest.param(
            replace(JUPYTER_MODE, python_cell_magics={"cust1", "cust2"}),
            "%%custom_python_magic -n1 -n2
x=2",
            pytest.raises(NothingChanged),
            id="No change when other cell magics registered",
        ),
        pytest.param(
            replace(JUPYTER_MODE, python_cell_magics={"custom_python_magic", "cust1"}),
            "%%custom_python_magic -n1 -n2
x = 2",
            does_not_raise(),
            id="Correctly change when cell magic registered",
        ),
    ],
)
def test_cell_magic_with_custom_python_magic(
    mode: Mode, expected_output: str, expectation: AbstractContextManager[object]
) -> None:
    with expectation:
        result = format_cell(
            "%%custom_python_magic -n1 -n2
x=2",
            fast=True,
            mode=mode,
        )
        assert result == expected_output


@pytest.mark.parametrize(
    "src",
    (
        "   %%custom_magic 
x=2",
        "

%%custom_magic
x=2",
        "# comment
%%custom_magic
x=2",
        "
  
 # comment with %%time
\t
 %%custom_magic # comment 
x=2",
    ),
)
def test_cell_magic_with_custom_python_magic_after_spaces_and_comments_noop(
    src: str,
) -> None:
    with pytest.raises(NothingChanged):
        format_cell(src, fast=True, mode=JUPYTER_MODE)


def test_cell_magic_nested() -> None:
    src = "%%time
%%time
2+2"
    result = format_cell(src, fast=True, mode=JUPYTER_MODE)
    expected = "%%time
%%time
2 + 2"
    assert result == expected


def test_cell_magic_with_magic_noop() -> None:
    src = "%%t -n1
ls = !ls"
    with pytest.raises(NothingChanged):
        format_cell(src, fast=True, mode=JUPYTER_MODE)


def test_automagic() -> None:
    src = "pip install black"
    with pytest.raises(NothingChanged):
        format_cell(src, fast=True, mode=JUPYTER_MODE)


def test_multiline_magic() -> None:
    src = "%time 1 + \\
2"
    with pytest.raises(NothingChanged):
        format_cell(src, fast=True, mode=JUPYTER_MODE)


def test_multiline_no_magic() -> None:
    src = "1 + \\
2"
    result = format_cell(src, fast=True, mode=JUPYTER_MODE)
    expected = "1 + 2"
    assert result == expected


def test_cell_magic_with_invalid_body() -> None:
    src = "%%time
if True"
    with pytest.raises(NothingChanged):
        format_cell(src, fast=True, mode=JUPYTER_MODE)


def test_empty_cell() -> None:
    src = ""
    with pytest.raises(NothingChanged):
        format_cell(src, fast=True, mode=JUPYTER_MODE)


def test_entire_notebook_empty_metadata() -> None:
    content = read_jupyter_notebook("jupyter", "notebook_empty_metadata")
    result = format_file_contents(content, fast=True, mode=JUPYTER_MODE)
    expected = (
        "{
"
        ' "cells": [
'
        "  {
"
        '   "cell_type": "code",
'
        '   "execution_count": null,
'
        '   "metadata": {
'
        '    "tags": []
'
        "   },
"
        '   "outputs": [],
'
        '   "source": [
'
        '    "%%time\
",
'
        '    "\
",
'
        '    "print(\\"foo\\")"
'
        "   ]
"
        "  },
"
        "  {
"
        '   "cell_type": "code",
'
        '   "execution_count": null,
'
        '   "metadata": {},
'
        '   "outputs": [],
'
        '   "source": []
'
        "  }
"
        " ],
"
        ' "metadata": {},
'
        ' "nbformat": 4,
'
        ' "nbformat_minor": 4
'
        "}
"
    )
    assert result == expected


def test_entire_notebook_trailing_newline() -> None:
    content = read_jupyter_notebook("jupyter", "notebook_trailing_newline")
    result = format_file_contents(content, fast=True, mode=JUPYTER_MODE)
    expected = (
        "{
"
        ' "cells": [
'
        "  {
"
        '   "cell_type": "code",
'
        '   "execution_count": null,
'
        '   "metadata": {
'
        '    "tags": []
'
        "   },
"
        '   "outputs": [],
'
        '   "source": [
'
        '    "%%time\
",
'
        '    "\
",
'
        '    "print(\\"foo\\")"
'
        "   ]
"
        "  },
"
        "  {
"
        '   "cell_type": "code",
'
        '   "execution_count": null,
'
        '   "metadata": {},
'
        '   "outputs": [],
'
        '   "source": []
'
        "  }
"
        " ],
"
        ' "metadata": {
'
        '  "interpreter": {
'
        '   "hash": "e758f3098b5b55f4d87fe30bbdc1367f20f246b483f96267ee70e6c40cb185d8"
'  # noqa:B950
        "  },
"
        '  "kernelspec": {
'
        '   "display_name": "Python 3.8.10 64-bit (\'black\': venv)",
'
        '   "name": "python3"
'
        "  },
"
        '  "language_info": {
'
        '   "name": "python",
'
        '   "version": ""
'
        "  }
"
        " },
"
        ' "nbformat": 4,
'
        ' "nbformat_minor": 4
'
        "}
"
    )
    assert result == expected


def test_entire_notebook_no_trailing_newline() -> None:
    content = read_jupyter_notebook("jupyter", "notebook_no_trailing_newline")
    result = format_file_contents(content, fast=True, mode=JUPYTER_MODE)
    expected = (
        "{
"
        ' "cells": [
'
        "  {
"
        '   "cell_type": "code",
'
        '   "execution_count": null,
'
        '   "metadata": {
'
        '    "tags": []
'
        "   },
"
        '   "outputs": [],
'
        '   "source": [
'
        '    "%%time\
",
'
        '    "\
",
'
        '    "print(\\"foo\\")"
'
        "   ]
"
        "  },
"
        "  {
"
        '   "cell_type": "code",
'
        '   "execution_count": null,
'
        '   "metadata": {},
'
        '   "outputs": [],
'
        '   "source": []
'
        "  }
"
        " ],
"
        ' "metadata": {
'
        '  "interpreter": {
'
        '   "hash": "e758f3098b5b55f4d87fe30bbdc1367f20f246b483f96267ee70e6c40cb185d8"
'  # noqa: B950
        "  },
"
        '  "kernelspec": {
'
        '   "display_name": "Python 3.8.10 64-bit (\'black\': venv)",
'
        '   "name": "python3"
'
        "  },
"
        '  "language_info": {
'
        '   "name": "python",
'
        '   "version": ""
'
        "  }
"
        " },
"
        ' "nbformat": 4,
'
        ' "nbformat_minor": 4
'
        "}"
    )
    assert result == expected


def test_entire_notebook_without_changes() -> None:
    content = read_jupyter_notebook("jupyter", "notebook_without_changes")
    with pytest.raises(NothingChanged):
        format_file_contents(content, fast=True, mode=JUPYTER_MODE)


def test_non_python_notebook() -> None:
    content = read_jupyter_notebook("jupyter", "non_python_notebook")

    with pytest.raises(NothingChanged):
        format_file_contents(content, fast=True, mode=JUPYTER_MODE)


def test_empty_string() -> None:
    with pytest.raises(NothingChanged):
        format_file_contents("", fast=True, mode=JUPYTER_MODE)


def test_unparseable_notebook() -> None:
    path = get_case_path("jupyter", "notebook_which_cant_be_parsed.ipynb")
    msg = rf"File '{re.escape(str(path))}' cannot be parsed as valid Jupyter notebook\."
    with pytest.raises(ValueError, match=msg):
        format_file_in_place(path, fast=True, mode=JUPYTER_MODE)


def test_ipynb_diff_with_change() -> None:
    result = runner.invoke(
        main,
        [
            str(get_case_path("jupyter", "notebook_trailing_newline.ipynb")),
            "--diff",
            f"--config={EMPTY_CONFIG}",
        ],
    )
    expected = "@@ -1,3 +1,3 @@
 %%time
 
-print('foo')
+print(\"foo\")
"
    assert expected in result.output


def test_ipynb_diff_with_no_change() -> None:
    result = runner.invoke(
        main,
        [
            str(get_case_path("jupyter", "notebook_without_changes.ipynb")),
            "--diff",
            f"--config={EMPTY_CONFIG}",
        ],
    )
    expected = "1 file would be left unchanged."
    assert expected in result.output


def test_cache_isnt_written_if_no_jupyter_deps_single(
    monkeypatch: MonkeyPatch, tmp_path: pathlib.Path
) -> None:
    # Check that the cache isn't written to if Jupyter dependencies aren't installed.
    jupyter_dependencies_are_installed.cache_clear()
    nb = get_case_path("jupyter", "notebook_trailing_newline.ipynb")
    tmp_nb = tmp_path / "notebook.ipynb"
    tmp_nb.write_bytes(nb.read_bytes())
    monkeypatch.setattr("black.jupyter_dependencies_are_installed", lambda warn: False)
    result = runner.invoke(
        main, [str(tmp_path / "notebook.ipynb"), f"--config={EMPTY_CONFIG}"]
    )
    assert "No Python files are present to be formatted. Nothing to do" in result.output
    jupyter_dependencies_are_installed.cache_clear()
    monkeypatch.setattr("black.jupyter_dependencies_are_installed", lambda warn: True)
    result = runner.invoke(
        main, [str(tmp_path / "notebook.ipynb"), f"--config={EMPTY_CONFIG}"]
    )
    assert "reformatted" in result.output


def test_cache_isnt_written_if_no_jupyter_deps_dir(
    monkeypatch: MonkeyPatch, tmp_path: pathlib.Path
) -> None:
    # Check that the cache isn't written to if Jupyter dependencies aren't installed.
    jupyter_dependencies_are_installed.cache_clear()
    nb = get_case_path("jupyter", "notebook_trailing_newline.ipynb")
    tmp_nb = tmp_path / "notebook.ipynb"
    tmp_nb.write_bytes(nb.read_bytes())
    monkeypatch.setattr(
        "black.files.jupyter_dependencies_are_installed", lambda warn: False
    )
    result = runner.invoke(main, [str(tmp_path), f"--config={EMPTY_CONFIG}"])
    assert "No Python files are present to be formatted. Nothing to do" in result.output
    jupyter_dependencies_are_installed.cache_clear()
    monkeypatch.setattr(
        "black.files.jupyter_dependencies_are_installed", lambda warn: True
    )
    result = runner.invoke(main, [str(tmp_path), f"--config={EMPTY_CONFIG}"])
    assert "reformatted" in result.output


def test_ipynb_flag(tmp_path: pathlib.Path) -> None:
    nb = get_case_path("jupyter", "notebook_trailing_newline.ipynb")
    tmp_nb = tmp_path / "notebook.a_file_extension_which_is_definitely_not_ipynb"
    tmp_nb.write_bytes(nb.read_bytes())
    result = runner.invoke(
        main,
        [
            str(tmp_nb),
            "--diff",
            "--ipynb",
            f"--config={EMPTY_CONFIG}",
        ],
    )
    expected = "@@ -1,3 +1,3 @@
 %%time
 
-print('foo')
+print(\"foo\")
"
    assert expected in result.output


def test_ipynb_and_pyi_flags() -> None:
    nb = get_case_path("jupyter", "notebook_trailing_newline.ipynb")
    result = runner.invoke(
        main,
        [
            str(nb),
            "--pyi",
            "--ipynb",
            "--diff",
            f"--config={EMPTY_CONFIG}",
        ],
    )
    assert isinstance(result.exception, SystemExit)
    expected = "Cannot pass both `pyi` and `ipynb` flags!
"
    assert result.output == expected


def test_unable_to_replace_magics(monkeypatch: MonkeyPatch) -> None:
    src = "%%time
a = 'foo'"
    monkeypatch.setattr("black.handle_ipynb_magics.TOKEN_HEX", lambda _: "foo")
    with pytest.raises(
        AssertionError, match="Black was not able to replace IPython magic"
    ):
        format_cell(src, fast=True, mode=JUPYTER_MODE)