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
  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     with open(DATA_DIR / "notebook_empty_metadata.ipynb", "rb") as fd:
 256         content_bytes = fd.read()
 257     content = content_bytes.decode()
 258     result = format_file_contents(content, fast=True, mode=JUPYTER_MODE)
 259     expected = (
 260         "{\n"
 261         ' "cells": [\n'
 262         "  {\n"
 263         '   "cell_type": "code",\n'
 264         '   "execution_count": null,\n'
 265         '   "metadata": {\n'
 266         '    "tags": []\n'
 267         "   },\n"
 268         '   "outputs": [],\n'
 269         '   "source": [\n'
 270         '    "%%time\\n",\n'
 271         '    "\\n",\n'
 272         '    "print(\\"foo\\")"\n'
 273         "   ]\n"
 274         "  },\n"
 275         "  {\n"
 276         '   "cell_type": "code",\n'
 277         '   "execution_count": null,\n'
 278         '   "metadata": {},\n'
 279         '   "outputs": [],\n'
 280         '   "source": []\n'
 281         "  }\n"
 282         " ],\n"
 283         ' "metadata": {},\n'
 284         ' "nbformat": 4,\n'
 285         ' "nbformat_minor": 4\n'
 286         "}\n"
 287     )
 288     assert result == expected
 289
 290
 291 def test_entire_notebook_trailing_newline() -> None:
 292     with open(DATA_DIR / "notebook_trailing_newline.ipynb", "rb") as fd:
 293         content_bytes = fd.read()
 294     content = content_bytes.decode()
 295     result = format_file_contents(content, fast=True, mode=JUPYTER_MODE)
 296     expected = (
 297         "{\n"
 298         ' "cells": [\n'
 299         "  {\n"
 300         '   "cell_type": "code",\n'
 301         '   "execution_count": null,\n'
 302         '   "metadata": {\n'
 303         '    "tags": []\n'
 304         "   },\n"
 305         '   "outputs": [],\n'
 306         '   "source": [\n'
 307         '    "%%time\\n",\n'
 308         '    "\\n",\n'
 309         '    "print(\\"foo\\")"\n'
 310         "   ]\n"
 311         "  },\n"
 312         "  {\n"
 313         '   "cell_type": "code",\n'
 314         '   "execution_count": null,\n'
 315         '   "metadata": {},\n'
 316         '   "outputs": [],\n'
 317         '   "source": []\n'
 318         "  }\n"
 319         " ],\n"
 320         ' "metadata": {\n'
 321         '  "interpreter": {\n'
 322         '   "hash": "e758f3098b5b55f4d87fe30bbdc1367f20f246b483f96267ee70e6c40cb185d8"\n'  # noqa:B950
 323         "  },\n"
 324         '  "kernelspec": {\n'
 325         '   "display_name": "Python 3.8.10 64-bit (\'black\': venv)",\n'
 326         '   "name": "python3"\n'
 327         "  },\n"
 328         '  "language_info": {\n'
 329         '   "name": "python",\n'
 330         '   "version": ""\n'
 331         "  }\n"
 332         " },\n"
 333         ' "nbformat": 4,\n'
 334         ' "nbformat_minor": 4\n'
 335         "}\n"
 336     )
 337     assert result == expected
 338
 339
 340 def test_entire_notebook_no_trailing_newline() -> None:
 341     with open(DATA_DIR / "notebook_no_trailing_newline.ipynb", "rb") as fd:
 342         content_bytes = fd.read()
 343     content = content_bytes.decode()
 344     result = format_file_contents(content, fast=True, mode=JUPYTER_MODE)
 345     expected = (
 346         "{\n"
 347         ' "cells": [\n'
 348         "  {\n"
 349         '   "cell_type": "code",\n'
 350         '   "execution_count": null,\n'
 351         '   "metadata": {\n'
 352         '    "tags": []\n'
 353         "   },\n"
 354         '   "outputs": [],\n'
 355         '   "source": [\n'
 356         '    "%%time\\n",\n'
 357         '    "\\n",\n'
 358         '    "print(\\"foo\\")"\n'
 359         "   ]\n"
 360         "  },\n"
 361         "  {\n"
 362         '   "cell_type": "code",\n'
 363         '   "execution_count": null,\n'
 364         '   "metadata": {},\n'
 365         '   "outputs": [],\n'
 366         '   "source": []\n'
 367         "  }\n"
 368         " ],\n"
 369         ' "metadata": {\n'
 370         '  "interpreter": {\n'
 371         '   "hash": "e758f3098b5b55f4d87fe30bbdc1367f20f246b483f96267ee70e6c40cb185d8"\n'  # noqa: B950
 372         "  },\n"
 373         '  "kernelspec": {\n'
 374         '   "display_name": "Python 3.8.10 64-bit (\'black\': venv)",\n'
 375         '   "name": "python3"\n'
 376         "  },\n"
 377         '  "language_info": {\n'
 378         '   "name": "python",\n'
 379         '   "version": ""\n'
 380         "  }\n"
 381         " },\n"
 382         ' "nbformat": 4,\n'
 383         ' "nbformat_minor": 4\n'
 384         "}"
 385     )
 386     assert result == expected
 387
 388
 389 def test_entire_notebook_without_changes() -> None:
 390     with open(DATA_DIR / "notebook_without_changes.ipynb", "rb") as fd:
 391         content_bytes = fd.read()
 392     content = content_bytes.decode()
 393     with pytest.raises(NothingChanged):
 394         format_file_contents(content, fast=True, mode=JUPYTER_MODE)
 395
 396
 397 def test_non_python_notebook() -> None:
 398     with open(DATA_DIR / "non_python_notebook.ipynb", "rb") as fd:
 399         content_bytes = fd.read()
 400     content = content_bytes.decode()
 401     with pytest.raises(NothingChanged):
 402         format_file_contents(content, fast=True, mode=JUPYTER_MODE)
 403
 404
 405 def test_empty_string() -> None:
 406     with pytest.raises(NothingChanged):
 407         format_file_contents("", fast=True, mode=JUPYTER_MODE)
 408
 409
 410 def test_unparseable_notebook() -> None:
 411     path = DATA_DIR / "notebook_which_cant_be_parsed.ipynb"
 412     msg = rf"File '{re.escape(str(path))}' cannot be parsed as valid Jupyter notebook\."
 413     with pytest.raises(ValueError, match=msg):
 414         format_file_in_place(path, fast=True, mode=JUPYTER_MODE)
 415
 416
 417 def test_ipynb_diff_with_change() -> None:
 418     result = runner.invoke(
 419         main,
 420         [
 421             str(DATA_DIR / "notebook_trailing_newline.ipynb"),
 422             "--diff",
 423             f"--config={EMPTY_CONFIG}",
 424         ],
 425     )
 426     expected = "@@ -1,3 +1,3 @@\n %%time\n \n-print('foo')\n+print(\"foo\")\n"
 427     assert expected in result.output
 428
 429
 430 def test_ipynb_diff_with_no_change() -> None:
 431     result = runner.invoke(
 432         main,
 433         [
 434             str(DATA_DIR / "notebook_without_changes.ipynb"),
 435             "--diff",
 436             f"--config={EMPTY_CONFIG}",
 437         ],
 438     )
 439     expected = "1 file would be left unchanged."
 440     assert expected in result.output
 441
 442
 443 def test_cache_isnt_written_if_no_jupyter_deps_single(
 444     monkeypatch: MonkeyPatch, tmp_path: pathlib.Path
 445 ) -> None:
 446     # Check that the cache isn't written to if Jupyter dependencies aren't installed.
 447     jupyter_dependencies_are_installed.cache_clear()
 448     nb = DATA_DIR / "notebook_trailing_newline.ipynb"
 449     tmp_nb = tmp_path / "notebook.ipynb"
 450     with open(nb) as src, open(tmp_nb, "w") as dst:
 451         dst.write(src.read())
 452     monkeypatch.setattr(
 453         "black.jupyter_dependencies_are_installed", lambda verbose, quiet: False
 454     )
 455     result = runner.invoke(
 456         main, [str(tmp_path / "notebook.ipynb"), f"--config={EMPTY_CONFIG}"]
 457     )
 458     assert "No Python files are present to be formatted. Nothing to do" in result.output
 459     jupyter_dependencies_are_installed.cache_clear()
 460     monkeypatch.setattr(
 461         "black.jupyter_dependencies_are_installed", lambda verbose, quiet: True
 462     )
 463     result = runner.invoke(
 464         main, [str(tmp_path / "notebook.ipynb"), f"--config={EMPTY_CONFIG}"]
 465     )
 466     assert "reformatted" in result.output
 467
 468
 469 def test_cache_isnt_written_if_no_jupyter_deps_dir(
 470     monkeypatch: MonkeyPatch, tmp_path: pathlib.Path
 471 ) -> None:
 472     # Check that the cache isn't written to if Jupyter dependencies aren't installed.
 473     jupyter_dependencies_are_installed.cache_clear()
 474     nb = DATA_DIR / "notebook_trailing_newline.ipynb"
 475     tmp_nb = tmp_path / "notebook.ipynb"
 476     with open(nb) as src, open(tmp_nb, "w") as dst:
 477         dst.write(src.read())
 478     monkeypatch.setattr(
 479         "black.files.jupyter_dependencies_are_installed", lambda verbose, quiet: False
 480     )
 481     result = runner.invoke(main, [str(tmp_path), f"--config={EMPTY_CONFIG}"])
 482     assert "No Python files are present to be formatted. Nothing to do" in result.output
 483     jupyter_dependencies_are_installed.cache_clear()
 484     monkeypatch.setattr(
 485         "black.files.jupyter_dependencies_are_installed", lambda verbose, quiet: True
 486     )
 487     result = runner.invoke(main, [str(tmp_path), f"--config={EMPTY_CONFIG}"])
 488     assert "reformatted" in result.output
 489
 490
 491 def test_ipynb_flag(tmp_path: pathlib.Path) -> None:
 492     nb = DATA_DIR / "notebook_trailing_newline.ipynb"
 493     tmp_nb = tmp_path / "notebook.a_file_extension_which_is_definitely_not_ipynb"
 494     with open(nb) as src, open(tmp_nb, "w") as dst:
 495         dst.write(src.read())
 496     result = runner.invoke(
 497         main,
 498         [
 499             str(tmp_nb),
 500             "--diff",
 501             "--ipynb",
 502             f"--config={EMPTY_CONFIG}",
 503         ],
 504     )
 505     expected = "@@ -1,3 +1,3 @@\n %%time\n \n-print('foo')\n+print(\"foo\")\n"
 506     assert expected in result.output
 507
 508
 509 def test_ipynb_and_pyi_flags() -> None:
 510     nb = DATA_DIR / "notebook_trailing_newline.ipynb"
 511     result = runner.invoke(
 512         main,
 513         [
 514             str(nb),
 515             "--pyi",
 516             "--ipynb",
 517             "--diff",
 518             f"--config={EMPTY_CONFIG}",
 519         ],
 520     )
 521     assert isinstance(result.exception, SystemExit)
 522     expected = "Cannot pass both `pyi` and `ipynb` flags!\n"
 523     assert result.output == expected
 524
 525
 526 def test_unable_to_replace_magics(monkeypatch: MonkeyPatch) -> None:
 527     src = "%%time\na = 'foo'"
 528     monkeypatch.setattr("black.handle_ipynb_magics.TOKEN_HEX", lambda _: "foo")
 529     with pytest.raises(
 530         AssertionError, match="Black was not able to replace IPython magic"
 531     ):
 532         format_cell(src, fast=True, mode=JUPYTER_MODE)