tests/test_ipynb.py

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