Repository: psf/black

Issue:
Allow defining custom python magics for Jupyter support
Recent updates of black have introduced out-of-the-box support for Jupyter Notebooks, which is extremely helpful.  As of #2630 only *known* python cell magics are processed (see e.g. [here](https://github.com/psf/black/blob/8a84bebcfcabddfd5b82a8cff0b830a745999b6c/src/black/handle_ipynb_magics.py#L40)), which certainly makes sense.  We have the use case that we have several custom python cell magics that we would also like to process.  

This request is to allow configuration to support additional (user-defined) magics.  


Possible solutions:

I am aware that we could certainly call black from python directly and programmatically update `PYTHON_CELL_MAGICS` to include our custom python magics.  However, I'm not keen on doing this because it would depend on implementation details within black.   If this would be the preferred solution, then I would suggest to at least add a function that could be called to register the additional magics.

Another possibility would be to add a configuration flag/option (i.e. in `pyproject.toml`/CLI) that would allow the input of a list of additional python cell magic names.  (Using an environment variable would be a similar option, but I don't see any other place that this is done, other than with e.g. `XDG_CACHE_HOME` which is a more "global" configuration variable.)  This is likely my preferred solution, but, given that black makes a concerted effort *to avoid* configuration options, I could understand why this would not be acceptable.

Thanks in advance for any input/consideration.



Discussion:
Hi @mgmarino ,

That's a good suggestion, thanks!

In nbQA I have such a configuration option ([cell magics](https://nbqa.readthedocs.io/en/latest/configuration.html#cell-magics)), but I left it out of the PR to `black` so as to not over-compliate things straight away.

Let's hear from the `black` maintainers if such a flag / configuration option would be welcome in `black` as well

> which is extremely helpful.

Glad to see this is useful!
Hi, thanks for submitting the issue! With every request like this I become more in favor of simplifying Black's responsibility and having a suite of plugins be implemented in other repositories as discussed in #2721.

However, given that we already have notebook support, this configuration would be really good. Perhaps named `python_cell_magics` like in our source?
@mgmarino do you want to submit a PR? [Here](https://github.com/psf/black/blob/main/CONTRIBUTING.md)'s the contributing guide if so
Sure, I'm happy to.  I'll have a look in the next few days.

Resolving patch:
diff --git a/src/black/__init__.py b/src/black/__init__.py
--- a/src/black/__init__.py
+++ b/src/black/__init__.py
@@ -24,6 +24,7 @@
     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
--- a/src/black/mode.py
+++ b/src/black/mode.py
@@ -4,6 +4,7 @@
 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 @@ def get_cache_key(self) -> str:
             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)


Test patch:
diff --git a/tests/test.toml b/tests/test.toml
--- 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
--- a/tests/test_black.py
+++ b/tests/test_black.py
@@ -1322,6 +1322,7 @@ def test_parse_pyproject_toml(self) -> None:
         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
--- 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)