tests/test_ipynb.py

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