From: Michael Marino Date: Fri, 21 Jan 2022 00:45:28 +0000 (+0100) Subject: Add support for custom python cell magics (#2744) X-Git-Url: https://git.madduck.net/etc/vim.git/commitdiff_plain/4ea75cd49521ed7fd8384e7a739e1abb1b6de46a Add support for custom python cell magics (#2744) Fixes #2742. This PR adds the ability to configure additional python cell magics. This will allow formatting cells in Jupyter Notebooks that are using custom (python) magics. --- diff --git a/CHANGES.md b/CHANGES.md index c3e2a33..a2e5c0a 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -30,6 +30,8 @@ - Fix handling of standalone `match()` or `case()` when there is a trailing newline or a comment inside of the parentheses. (#2760) - Black now normalizes string prefix order (#2297) +- Add configuration option (`python-cell-magics`) to format cells with custom magics in + Jupyter Notebooks (#2744) - Deprecate `--experimental-string-processing` and move the functionality under `--preview` (#2789) diff --git a/src/black/__init__.py b/src/black/__init__.py index bdece68..eaf72f9 100644 --- a/src/black/__init__.py +++ b/src/black/__init__.py @@ -24,6 +24,7 @@ from typing import ( MutableMapping, Optional, Pattern, + Sequence, Set, Sized, Tuple, @@ -225,6 +226,16 @@ def validate_regex( "(useful when piping source on standard input)." ), ) +@click.option( + "--python-cell-magics", + multiple=True, + help=( + "When processing Jupyter Notebooks, add the given magic to the list" + f" of known python-magics ({', '.join(PYTHON_CELL_MAGICS)})." + " Useful for formatting cells with custom python magics." + ), + default=[], +) @click.option( "-S", "--skip-string-normalization", @@ -401,6 +412,7 @@ def main( fast: bool, pyi: bool, ipynb: bool, + python_cell_magics: Sequence[str], skip_string_normalization: bool, skip_magic_trailing_comma: bool, experimental_string_processing: bool, @@ -476,6 +488,7 @@ def main( magic_trailing_comma=not skip_magic_trailing_comma, experimental_string_processing=experimental_string_processing, preview=preview, + python_cell_magics=set(python_cell_magics), ) if code is not None: @@ -981,7 +994,7 @@ def format_file_contents(src_contents: str, *, fast: bool, mode: Mode) -> FileCo return dst_contents -def validate_cell(src: str) -> None: +def validate_cell(src: str, mode: Mode) -> None: """Check that cell does not already contain TransformerManager transformations, or non-Python cell magics, which might cause tokenizer_rt to break because of indentations. @@ -1000,7 +1013,10 @@ def validate_cell(src: str) -> None: """ if any(transformed_magic in src for transformed_magic in TRANSFORMED_MAGICS): raise NothingChanged - if src[:2] == "%%" and src.split()[0][2:] not in PYTHON_CELL_MAGICS: + if ( + src[:2] == "%%" + and src.split()[0][2:] not in PYTHON_CELL_MAGICS | mode.python_cell_magics + ): raise NothingChanged @@ -1020,7 +1036,7 @@ def format_cell(src: str, *, fast: bool, mode: Mode) -> str: could potentially be automagics or multi-line magics, which are currently not supported. """ - validate_cell(src) + validate_cell(src, mode) src_without_trailing_semicolon, has_trailing_semicolon = remove_trailing_semicolon( src ) diff --git a/src/black/mode.py b/src/black/mode.py index b6d1a1f..6d45e3d 100644 --- a/src/black/mode.py +++ b/src/black/mode.py @@ -4,6 +4,7 @@ Mostly around Python language feature support per version and Black configuratio chosen by the user. """ +from hashlib import md5 import sys from dataclasses import dataclass, field @@ -142,6 +143,7 @@ class Mode: is_ipynb: bool = False magic_trailing_comma: bool = True experimental_string_processing: bool = False + python_cell_magics: Set[str] = field(default_factory=set) preview: bool = False def __post_init__(self) -> None: @@ -180,5 +182,6 @@ class Mode: str(int(self.magic_trailing_comma)), str(int(self.experimental_string_processing)), str(int(self.preview)), + md5((",".join(sorted(self.python_cell_magics))).encode()).hexdigest(), ] return ".".join(parts) diff --git a/tests/test.toml b/tests/test.toml index d3ab1e6..e5fb922 100644 --- a/tests/test.toml +++ b/tests/test.toml @@ -7,6 +7,7 @@ line-length = 79 target-version = ["py36", "py37", "py38"] exclude='\.pyi?$' include='\.py?$' +python-cell-magics = ["custom1", "custom2"] [v1.0.0-syntax] # This shouldn't break Black. diff --git a/tests/test_black.py b/tests/test_black.py index 19cff23..fd01425 100644 --- a/tests/test_black.py +++ b/tests/test_black.py @@ -1322,6 +1322,7 @@ class BlackTestCase(BlackBaseTestCase): self.assertEqual(config["color"], True) self.assertEqual(config["line_length"], 79) self.assertEqual(config["target_version"], ["py36", "py37", "py38"]) + self.assertEqual(config["python_cell_magics"], ["custom1", "custom2"]) self.assertEqual(config["exclude"], r"\.pyi?$") self.assertEqual(config["include"], r"\.py?$") diff --git a/tests/test_ipynb.py b/tests/test_ipynb.py index fe8d67a..d78a68c 100644 --- a/tests/test_ipynb.py +++ b/tests/test_ipynb.py @@ -1,5 +1,8 @@ +from dataclasses import replace import pathlib import re +from contextlib import ExitStack as does_not_raise +from typing import ContextManager from click.testing import CliRunner from black.handle_ipynb_magics import jupyter_dependencies_are_installed @@ -63,9 +66,19 @@ def test_trailing_semicolon_noop() -> None: format_cell(src, fast=True, mode=JUPYTER_MODE) -def test_cell_magic() -> None: +@pytest.mark.parametrize( + "mode", + [ + pytest.param(JUPYTER_MODE, id="default mode"), + pytest.param( + replace(JUPYTER_MODE, python_cell_magics={"cust1", "cust1"}), + id="custom cell magics mode", + ), + ], +) +def test_cell_magic(mode: Mode) -> None: src = "%%time\nfoo =bar" - result = format_cell(src, fast=True, mode=JUPYTER_MODE) + result = format_cell(src, fast=True, mode=mode) expected = "%%time\nfoo = bar" assert result == expected @@ -76,6 +89,16 @@ def test_cell_magic_noop() -> None: format_cell(src, fast=True, mode=JUPYTER_MODE) +@pytest.mark.parametrize( + "mode", + [ + pytest.param(JUPYTER_MODE, id="default mode"), + pytest.param( + replace(JUPYTER_MODE, python_cell_magics={"cust1", "cust1"}), + id="custom cell magics mode", + ), + ], +) @pytest.mark.parametrize( "src, expected", ( @@ -96,8 +119,8 @@ def test_cell_magic_noop() -> None: pytest.param("env = %env", "env = %env", id="Assignment to magic"), ), ) -def test_magic(src: str, expected: str) -> None: - result = format_cell(src, fast=True, mode=JUPYTER_MODE) +def test_magic(src: str, expected: str, mode: Mode) -> None: + result = format_cell(src, fast=True, mode=mode) assert result == expected @@ -139,6 +162,41 @@ def test_cell_magic_with_magic() -> None: assert result == expected +@pytest.mark.parametrize( + "mode, expected_output, expectation", + [ + pytest.param( + JUPYTER_MODE, + "%%custom_python_magic -n1 -n2\nx=2", + pytest.raises(NothingChanged), + id="No change when cell magic not registered", + ), + pytest.param( + replace(JUPYTER_MODE, python_cell_magics={"cust1", "cust1"}), + "%%custom_python_magic -n1 -n2\nx=2", + pytest.raises(NothingChanged), + id="No change when other cell magics registered", + ), + pytest.param( + replace(JUPYTER_MODE, python_cell_magics={"custom_python_magic", "cust1"}), + "%%custom_python_magic -n1 -n2\nx = 2", + does_not_raise(), + id="Correctly change when cell magic registered", + ), + ], +) +def test_cell_magic_with_custom_python_magic( + mode: Mode, expected_output: str, expectation: ContextManager[object] +) -> None: + with expectation: + result = format_cell( + "%%custom_python_magic -n1 -n2\nx=2", + fast=True, + mode=mode, + ) + assert result == expected_output + + def test_cell_magic_nested() -> None: src = "%%time\n%%time\n2+2" result = format_cell(src, fast=True, mode=JUPYTER_MODE)