Add support for custom python cell magics

CHANGES.md

@@ -17,6 +17,8 @@
 - Tuple unpacking on `return` and `yield` constructs now implies 3.8+ (#2700)
 - Unparenthesized tuples on annotated assignments (e.g
   `values: Tuple[int, ...] = 1, 2, 3`) now implies 3.8+ (#2708)
+- Add configuration option (`python-cell-magics`) to format cells with custom magics in
+  Jupyter Notebooks (#2744)
 
 ### Packaging
 

src/black/__init__.py

@@ -224,6 +224,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",
@@ -395,6 +405,7 @@ def main(
     fast: bool,
     pyi: bool,
     ipynb: bool,
+    python_cell_magics: List[str],
     skip_string_normalization: bool,
     skip_magic_trailing_comma: bool,
     experimental_string_processing: bool,
@@ -439,6 +450,7 @@ def main(
         string_normalization=not skip_string_normalization,
         magic_trailing_comma=not skip_magic_trailing_comma,
         experimental_string_processing=experimental_string_processing,
+        python_cell_magics=set(python_cell_magics),
     )
 
     if code is not None:
@@ -944,7 +956,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.
@@ -963,7 +975,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
 
 
@@ -983,7 +998,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
     )

src/black/mode.py

@@ -166,6 +166,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)
 
     def get_cache_key(self) -> str:
         if self.target_versions:
@@ -183,5 +184,6 @@ def get_cache_key(self) -> str:
             str(int(self.is_ipynb)),
             str(int(self.magic_trailing_comma)),
             str(int(self.experimental_string_processing)),
+            ",".join(self.python_cell_magics),
         ]
         return ".".join(parts)

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.

tests/test_black.py

@@ -1324,6 +1324,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?$")
 

tests/test_ipynb.py

@@ -1,3 +1,4 @@
+from dataclasses import replace
 import pathlib
 import re
 
@@ -139,6 +140,17 @@ def test_cell_magic_with_magic() -> None:
     assert result == expected
 
 
+def test_cell_magic_with_custom_python_magic() -> None:
+    src = "%%custom_python_magic -n1 -n2\nx=2"
+    result = format_cell(
+        src,
+        fast=True,
+        mode=replace(JUPYTER_MODE, python_cell_magics={"custom_python_magic"}),
+    )
+    expected = "%%custom_python_magic -n1 -n2\nx = 2"
+    assert result == expected
+
+
 def test_cell_magic_nested() -> None:
     src = "%%time\n%%time\n2+2"
     result = format_cell(src, fast=True, mode=JUPYTER_MODE)