tests/test_ipynb.py

   1 import contextlib
   2 from dataclasses import replace
   3 import pathlib
   4 import re
   5 from contextlib import ExitStack as does_not_raise
   6 from typing import ContextManager
   7
   8 from click.testing import CliRunner
   9 from black.handle_ipynb_magics import jupyter_dependencies_are_installed
  10 from black import (
  11     main,
  12     NothingChanged,
  13     format_cell,
  14     format_file_contents,
  15     format_file_in_place,
  16 )
  17 import pytest
  18 from black import Mode
  19 from _pytest.monkeypatch import MonkeyPatch
  20 from tests.util import DATA_DIR, read_jupyter_notebook, get_case_path
  21
  22 with contextlib.suppress(ModuleNotFoundError):
  23     import IPython
  24 pytestmark = pytest.mark.jupyter
  25 pytest.importorskip("IPython", reason="IPython is an optional dependency")
  26 pytest.importorskip("tokenize_rt", reason="tokenize-rt is an optional dependency")
  27
  28 JUPYTER_MODE = Mode(is_ipynb=True)
  29
  30 EMPTY_CONFIG = DATA_DIR / "empty_pyproject.toml"
  31
  32 runner = CliRunner()
  33
  34
  35 def test_noop() -> None:
  36     src = 'foo = "a"'
  37     with pytest.raises(NothingChanged):
  38         format_cell(src, fast=True, mode=JUPYTER_MODE)
  39
  40
  41 @pytest.mark.parametrize("fast", [True, False])
  42 def test_trailing_semicolon(fast: bool) -> None:
  43     src = 'foo = "a" ;'
  44     result = format_cell(src, fast=fast, mode=JUPYTER_MODE)
  45     expected = 'foo = "a";'
  46     assert result == expected
  47
  48
  49 def test_trailing_semicolon_with_comment() -> None:
  50     src = 'foo = "a" ;  # bar'
  51     result = format_cell(src, fast=True, mode=JUPYTER_MODE)
  52     expected = 'foo = "a";  # bar'
  53     assert result == expected
  54
  55
  56 def test_trailing_semicolon_with_comment_on_next_line() -> None:
  57     src = "import black;\n\n# this is a comment"
  58     with pytest.raises(NothingChanged):
  59         format_cell(src, fast=True, mode=JUPYTER_MODE)
  60
  61
  62 def test_trailing_semicolon_indented() -> None:
  63     src = "with foo:\n    plot_bar();"
  64     with pytest.raises(NothingChanged):
  65         format_cell(src, fast=True, mode=JUPYTER_MODE)
  66
  67
  68 def test_trailing_semicolon_noop() -> None:
  69     src = 'foo = "a";'
  70     with pytest.raises(NothingChanged):
  71         format_cell(src, fast=True, mode=JUPYTER_MODE)
  72
  73
  74 @pytest.mark.parametrize(
  75     "mode",
  76     [
  77         pytest.param(JUPYTER_MODE, id="default mode"),
  78         pytest.param(
  79             replace(JUPYTER_MODE, python_cell_magics={"cust1", "cust1"}),
  80             id="custom cell magics mode",
  81         ),
  82     ],
  83 )
  84 def test_cell_magic(mode: Mode) -> None:
  85     src = "%%time\nfoo =bar"
  86     result = format_cell(src, fast=True, mode=mode)
  87     expected = "%%time\nfoo = bar"
  88     assert result == expected
  89
  90
  91 def test_cell_magic_noop() -> None:
  92     src = "%%time\n2 + 2"
  93     with pytest.raises(NothingChanged):
  94         format_cell(src, fast=True, mode=JUPYTER_MODE)
  95
  96
  97 @pytest.mark.parametrize(
  98     "mode",
  99     [
 100         pytest.param(JUPYTER_MODE, id="default mode"),
 101         pytest.param(
 102             replace(JUPYTER_MODE, python_cell_magics={"cust1", "cust1"}),
 103             id="custom cell magics mode",
 104         ),
 105     ],
 106 )
 107 @pytest.mark.parametrize(
 108     "src, expected",
 109     (
 110         pytest.param("ls =!ls", "ls = !ls", id="System assignment"),
 111         pytest.param("!ls\n'foo'", '!ls\n"foo"', id="System call"),
 112         pytest.param("!!ls\n'foo'", '!!ls\n"foo"', id="Other system call"),
 113         pytest.param("?str\n'foo'", '?str\n"foo"', id="Help"),
 114         pytest.param("??str\n'foo'", '??str\n"foo"', id="Other help"),
 115         pytest.param(
 116             "%matplotlib inline\n'foo'",
 117             '%matplotlib inline\n"foo"',
 118             id="Line magic with argument",
 119         ),
 120         pytest.param("%time\n'foo'", '%time\n"foo"', id="Line magic without argument"),
 121         pytest.param(
 122             "env =  %env var", "env = %env var", id="Assignment to environment variable"
 123         ),
 124         pytest.param("env =  %env", "env = %env", id="Assignment to magic"),
 125     ),
 126 )
 127 def test_magic(src: str, expected: str, mode: Mode) -> None:
 128     result = format_cell(src, fast=True, mode=mode)
 129     assert result == expected
 130
 131
 132 @pytest.mark.parametrize(
 133     "src",
 134     (
 135         "%%bash\n2+2",
 136         "%%html --isolated\n2+2",
 137         "%%writefile e.txt\n  meh\n meh",
 138     ),
 139 )
 140 def test_non_python_magics(src: str) -> None:
 141     with pytest.raises(NothingChanged):
 142         format_cell(src, fast=True, mode=JUPYTER_MODE)
 143
 144
 145 @pytest.mark.skipif(
 146     IPython.version_info < (8, 3),
 147     reason="Change in how TransformerManager transforms this input",
 148 )
 149 def test_set_input() -> None:
 150     src = "a = b??"
 151     expected = "??b"
 152     result = format_cell(src, fast=True, mode=JUPYTER_MODE)
 153     assert result == expected
 154
 155
 156 def test_input_already_contains_transformed_magic() -> None:
 157     src = '%time foo()\nget_ipython().run_cell_magic("time", "", "foo()\\n")'
 158     with pytest.raises(NothingChanged):
 159         format_cell(src, fast=True, mode=JUPYTER_MODE)
 160
 161
 162 def test_magic_noop() -> None:
 163     src = "ls = !ls"
 164     with pytest.raises(NothingChanged):
 165         format_cell(src, fast=True, mode=JUPYTER_MODE)
 166
 167
 168 def test_cell_magic_with_magic() -> None:
 169     src = "%%timeit -n1\nls =!ls"
 170     result = format_cell(src, fast=True, mode=JUPYTER_MODE)
 171     expected = "%%timeit -n1\nls = !ls"
 172     assert result == expected
 173
 174
 175 @pytest.mark.parametrize(
 176     "mode, expected_output, expectation",
 177     [
 178         pytest.param(
 179             JUPYTER_MODE,
 180             "%%custom_python_magic -n1 -n2\nx=2",
 181             pytest.raises(NothingChanged),
 182             id="No change when cell magic not registered",
 183         ),
 184         pytest.param(
 185             replace(JUPYTER_MODE, python_cell_magics={"cust1", "cust1"}),
 186             "%%custom_python_magic -n1 -n2\nx=2",
 187             pytest.raises(NothingChanged),
 188             id="No change when other cell magics registered",
 189         ),
 190         pytest.param(
 191             replace(JUPYTER_MODE, python_cell_magics={"custom_python_magic", "cust1"}),
 192             "%%custom_python_magic -n1 -n2\nx = 2",
 193             does_not_raise(),
 194             id="Correctly change when cell magic registered",
 195         ),
 196     ],
 197 )
 198 def test_cell_magic_with_custom_python_magic(
 199     mode: Mode, expected_output: str, expectation: ContextManager[object]
 200 ) -> None:
 201     with expectation:
 202         result = format_cell(
 203             "%%custom_python_magic -n1 -n2\nx=2",
 204             fast=True,
 205             mode=mode,
 206         )
 207         assert result == expected_output
 208
 209
 210 def test_cell_magic_nested() -> None:
 211     src = "%%time\n%%time\n2+2"
 212     result = format_cell(src, fast=True, mode=JUPYTER_MODE)
 213     expected = "%%time\n%%time\n2 + 2"
 214     assert result == expected
 215
 216
 217 def test_cell_magic_with_magic_noop() -> None:
 218     src = "%%t -n1\nls = !ls"
 219     with pytest.raises(NothingChanged):
 220         format_cell(src, fast=True, mode=JUPYTER_MODE)
 221
 222
 223 def test_automagic() -> None:
 224     src = "pip install black"
 225     with pytest.raises(NothingChanged):
 226         format_cell(src, fast=True, mode=JUPYTER_MODE)
 227
 228
 229 def test_multiline_magic() -> None:
 230     src = "%time 1 + \\\n2"
 231     with pytest.raises(NothingChanged):
 232         format_cell(src, fast=True, mode=JUPYTER_MODE)
 233
 234
 235 def test_multiline_no_magic() -> None:
 236     src = "1 + \\\n2"
 237     result = format_cell(src, fast=True, mode=JUPYTER_MODE)
 238     expected = "1 + 2"
 239     assert result == expected
 240
 241
 242 def test_cell_magic_with_invalid_body() -> None:
 243     src = "%%time\nif True"
 244     with pytest.raises(NothingChanged):
 245         format_cell(src, fast=True, mode=JUPYTER_MODE)
 246
 247
 248 def test_empty_cell() -> None:
 249     src = ""
 250     with pytest.raises(NothingChanged):
 251         format_cell(src, fast=True, mode=JUPYTER_MODE)
 252
 253
 254 def test_entire_notebook_empty_metadata() -> None:
 255     content = read_jupyter_notebook("jupyter", "notebook_empty_metadata")
 256     result = format_file_contents(content, fast=True, mode=JUPYTER_MODE)
 257     expected = (
 258         "{\n"
 259         ' "cells": [\n'
 260         "  {\n"
 261         '   "cell_type": "code",\n'
 262         '   "execution_count": null,\n'
 263         '   "metadata": {\n'
 264         '    "tags": []\n'
 265         "   },\n"
 266         '   "outputs": [],\n'
 267         '   "source": [\n'
 268         '    "%%time\\n",\n'
 269         '    "\\n",\n'
 270         '    "print(\\"foo\\")"\n'
 271         "   ]\n"
 272         "  },\n"
 273         "  {\n"
 274         '   "cell_type": "code",\n'
 275         '   "execution_count": null,\n'
 276         '   "metadata": {},\n'
 277         '   "outputs": [],\n'
 278         '   "source": []\n'
 279         "  }\n"
 280         " ],\n"
 281         ' "metadata": {},\n'
 282         ' "nbformat": 4,\n'
 283         ' "nbformat_minor": 4\n'
 284         "}\n"
 285     )
 286     assert result == expected
 287
 288
 289 def test_entire_notebook_trailing_newline() -> None:
 290     content = read_jupyter_notebook("jupyter", "notebook_trailing_newline")
 291     result = format_file_contents(content, fast=True, mode=JUPYTER_MODE)
 292     expected = (
 293         "{\n"
 294         ' "cells": [\n'
 295         "  {\n"
 296         '   "cell_type": "code",\n'
 297         '   "execution_count": null,\n'
 298         '   "metadata": {\n'
 299         '    "tags": []\n'
 300         "   },\n"
 301         '   "outputs": [],\n'
 302         '   "source": [\n'
 303         '    "%%time\\n",\n'
 304         '    "\\n",\n'
 305         '    "print(\\"foo\\")"\n'
 306         "   ]\n"
 307         "  },\n"
 308         "  {\n"
 309         '   "cell_type": "code",\n'
 310         '   "execution_count": null,\n'
 311         '   "metadata": {},\n'
 312         '   "outputs": [],\n'
 313         '   "source": []\n'
 314         "  }\n"
 315         " ],\n"
 316         ' "metadata": {\n'
 317         '  "interpreter": {\n'
 318         '   "hash": "e758f3098b5b55f4d87fe30bbdc1367f20f246b483f96267ee70e6c40cb185d8"\n'  # noqa:B950
 319         "  },\n"
 320         '  "kernelspec": {\n'
 321         '   "display_name": "Python 3.8.10 64-bit (\'black\': venv)",\n'
 322         '   "name": "python3"\n'
 323         "  },\n"
 324         '  "language_info": {\n'
 325         '   "name": "python",\n'
 326         '   "version": ""\n'
 327         "  }\n"
 328         " },\n"
 329         ' "nbformat": 4,\n'
 330         ' "nbformat_minor": 4\n'
 331         "}\n"
 332     )
 333     assert result == expected
 334
 335
 336 def test_entire_notebook_no_trailing_newline() -> None:
 337     content = read_jupyter_notebook("jupyter", "notebook_no_trailing_newline")
 338     result = format_file_contents(content, fast=True, mode=JUPYTER_MODE)
 339     expected = (
 340         "{\n"
 341         ' "cells": [\n'
 342         "  {\n"
 343         '   "cell_type": "code",\n'
 344         '   "execution_count": null,\n'
 345         '   "metadata": {\n'
 346         '    "tags": []\n'
 347         "   },\n"
 348         '   "outputs": [],\n'
 349         '   "source": [\n'
 350         '    "%%time\\n",\n'
 351         '    "\\n",\n'
 352         '    "print(\\"foo\\")"\n'
 353         "   ]\n"
 354         "  },\n"
 355         "  {\n"
 356         '   "cell_type": "code",\n'
 357         '   "execution_count": null,\n'
 358         '   "metadata": {},\n'
 359         '   "outputs": [],\n'
 360         '   "source": []\n'
 361         "  }\n"
 362         " ],\n"
 363         ' "metadata": {\n'
 364         '  "interpreter": {\n'
 365         '   "hash": "e758f3098b5b55f4d87fe30bbdc1367f20f246b483f96267ee70e6c40cb185d8"\n'  # noqa: B950
 366         "  },\n"
 367         '  "kernelspec": {\n'
 368         '   "display_name": "Python 3.8.10 64-bit (\'black\': venv)",\n'
 369         '   "name": "python3"\n'
 370         "  },\n"
 371         '  "language_info": {\n'
 372         '   "name": "python",\n'
 373         '   "version": ""\n'
 374         "  }\n"
 375         " },\n"
 376         ' "nbformat": 4,\n'
 377         ' "nbformat_minor": 4\n'
 378         "}"
 379     )
 380     assert result == expected
 381
 382
 383 def test_entire_notebook_without_changes() -> None:
 384     content = read_jupyter_notebook("jupyter", "notebook_without_changes")
 385     with pytest.raises(NothingChanged):
 386         format_file_contents(content, fast=True, mode=JUPYTER_MODE)
 387
 388
 389 def test_non_python_notebook() -> None:
 390     content = read_jupyter_notebook("jupyter", "non_python_notebook")
 391
 392     with pytest.raises(NothingChanged):
 393         format_file_contents(content, fast=True, mode=JUPYTER_MODE)
 394
 395
 396 def test_empty_string() -> None:
 397     with pytest.raises(NothingChanged):
 398         format_file_contents("", fast=True, mode=JUPYTER_MODE)
 399
 400
 401 def test_unparseable_notebook() -> None:
 402     path = get_case_path("jupyter", "notebook_which_cant_be_parsed.ipynb")
 403     msg = rf"File '{re.escape(str(path))}' cannot be parsed as valid Jupyter notebook\."
 404     with pytest.raises(ValueError, match=msg):
 405         format_file_in_place(path, fast=True, mode=JUPYTER_MODE)
 406
 407
 408 def test_ipynb_diff_with_change() -> None:
 409     result = runner.invoke(
 410         main,
 411         [
 412             str(get_case_path("jupyter", "notebook_trailing_newline.ipynb")),
 413             "--diff",
 414             f"--config={EMPTY_CONFIG}",
 415         ],
 416     )
 417     expected = "@@ -1,3 +1,3 @@\n %%time\n \n-print('foo')\n+print(\"foo\")\n"
 418     assert expected in result.output
 419
 420
 421 def test_ipynb_diff_with_no_change() -> None:
 422     result = runner.invoke(
 423         main,
 424         [
 425             str(get_case_path("jupyter", "notebook_without_changes.ipynb")),
 426             "--diff",
 427             f"--config={EMPTY_CONFIG}",
 428         ],
 429     )
 430     expected = "1 file would be left unchanged."
 431     assert expected in result.output
 432
 433
 434 def test_cache_isnt_written_if_no_jupyter_deps_single(
 435     monkeypatch: MonkeyPatch, tmp_path: pathlib.Path
 436 ) -> None:
 437     # Check that the cache isn't written to if Jupyter dependencies aren't installed.
 438     jupyter_dependencies_are_installed.cache_clear()
 439     nb = get_case_path("jupyter", "notebook_trailing_newline.ipynb")
 440     tmp_nb = tmp_path / "notebook.ipynb"
 441     with open(nb) as src, open(tmp_nb, "w") as dst:
 442         dst.write(src.read())
 443     monkeypatch.setattr(
 444         "black.jupyter_dependencies_are_installed", lambda verbose, quiet: False
 445     )
 446     result = runner.invoke(
 447         main, [str(tmp_path / "notebook.ipynb"), f"--config={EMPTY_CONFIG}"]
 448     )
 449     assert "No Python files are present to be formatted. Nothing to do" in result.output
 450     jupyter_dependencies_are_installed.cache_clear()
 451     monkeypatch.setattr(
 452         "black.jupyter_dependencies_are_installed", lambda verbose, quiet: True
 453     )
 454     result = runner.invoke(
 455         main, [str(tmp_path / "notebook.ipynb"), f"--config={EMPTY_CONFIG}"]
 456     )
 457     assert "reformatted" in result.output
 458
 459
 460 def test_cache_isnt_written_if_no_jupyter_deps_dir(
 461     monkeypatch: MonkeyPatch, tmp_path: pathlib.Path
 462 ) -> None:
 463     # Check that the cache isn't written to if Jupyter dependencies aren't installed.
 464     jupyter_dependencies_are_installed.cache_clear()
 465     nb = get_case_path("jupyter", "notebook_trailing_newline.ipynb")
 466     tmp_nb = tmp_path / "notebook.ipynb"
 467     with open(nb) as src, open(tmp_nb, "w") as dst:
 468         dst.write(src.read())
 469     monkeypatch.setattr(
 470         "black.files.jupyter_dependencies_are_installed", lambda verbose, quiet: False
 471     )
 472     result = runner.invoke(main, [str(tmp_path), f"--config={EMPTY_CONFIG}"])
 473     assert "No Python files are present to be formatted. Nothing to do" in result.output
 474     jupyter_dependencies_are_installed.cache_clear()
 475     monkeypatch.setattr(
 476         "black.files.jupyter_dependencies_are_installed", lambda verbose, quiet: True
 477     )
 478     result = runner.invoke(main, [str(tmp_path), f"--config={EMPTY_CONFIG}"])
 479     assert "reformatted" in result.output
 480
 481
 482 def test_ipynb_flag(tmp_path: pathlib.Path) -> None:
 483     nb = get_case_path("jupyter", "notebook_trailing_newline.ipynb")
 484     tmp_nb = tmp_path / "notebook.a_file_extension_which_is_definitely_not_ipynb"
 485     with open(nb) as src, open(tmp_nb, "w") as dst:
 486         dst.write(src.read())
 487     result = runner.invoke(
 488         main,
 489         [
 490             str(tmp_nb),
 491             "--diff",
 492             "--ipynb",
 493             f"--config={EMPTY_CONFIG}",
 494         ],
 495     )
 496     expected = "@@ -1,3 +1,3 @@\n %%time\n \n-print('foo')\n+print(\"foo\")\n"
 497     assert expected in result.output
 498
 499
 500 def test_ipynb_and_pyi_flags() -> None:
 501     nb = get_case_path("jupyter", "notebook_trailing_newline.ipynb")
 502     result = runner.invoke(
 503         main,
 504         [
 505             str(nb),
 506             "--pyi",
 507             "--ipynb",
 508             "--diff",
 509             f"--config={EMPTY_CONFIG}",
 510         ],
 511     )
 512     assert isinstance(result.exception, SystemExit)
 513     expected = "Cannot pass both `pyi` and `ipynb` flags!\n"
 514     assert result.output == expected
 515
 516
 517 def test_unable_to_replace_magics(monkeypatch: MonkeyPatch) -> None:
 518     src = "%%time\na = 'foo'"
 519     monkeypatch.setattr("black.handle_ipynb_magics.TOKEN_HEX", lambda _: "foo")
 520     with pytest.raises(
 521         AssertionError, match="Black was not able to replace IPython magic"
 522     ):
 523         format_cell(src, fast=True, mode=JUPYTER_MODE)