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