Add support for on/off comments

README.rst

@@ -119,6 +119,37 @@ It also has the below extra options:
 * ``-E`` / ``--skip-errors`` - Don’t exit non-zero for errors from Black (normally syntax errors).
 * ``--rst-literal-blocks`` - Also format literal blocks in reStructuredText files (more below).
 
+To prevent formatting in specific regions of a document,
+use comments to disable/re-enable formatting:
+
+.. code-block:: markdown
+
+    <!-- blacken-docs:off -->
+    ```python
+    f(1, 2, 3)
+    ```
+    <!-- blacken-docs:on -->
+
+.. code-block:: rst
+
+    ..
+        blacken-docs:off
+
+    .. code-block:: python
+
+        f(1,2,3)
+
+    ..
+        blacken-docs:on
+
+.. code-block:: latex
+
+    % blacken-docs:off
+    \begin{minted}{python}
+    f(1, 2, 3)
+    \end{minted}
+    % blacken-docs:on
+
 History
 =======
 

src/blacken_docs/__init__.py

@@ -86,6 +86,13 @@
 )
 INDENT_RE = re.compile("^ +(?=[^ ])", re.MULTILINE)
 TRAILING_NL_RE = re.compile(r"\n+\Z", re.MULTILINE)
+ON_OFF = r"blacken-docs:\s*(on|off)"
+ON_OFF_COMMENT_RE = re.compile(
+    rf"(?:^\s*<!--\s+{ON_OFF}\s+-->\s*$)|"
+    rf"(?:^(?P<indent>\s*)\.\.\n(?P=indent) +{ON_OFF}\s*$)|"
+    rf"(?:^\s*%\s*{ON_OFF}\s*$)",
+    re.MULTILINE,
+)
 
 
 class CodeBlockError:
@@ -102,6 +109,26 @@ def format_str(
 ) -> tuple[str, Sequence[CodeBlockError]]:
     errors: list[CodeBlockError] = []
 
+    off_ranges = []
+    off_start = None
+    for comment in re.finditer(ON_OFF_COMMENT_RE, src):
+        # In the `ON_OFF_COMMENT_RE` regex, we cannot use the same group name
+        # multiple times, and group numbers are not reset from one alternative
+        # to another (`r"(alt1|alt2|...)`), so we rely on the fact that the
+        # `(on|off)` group always appear last in each alternative, the other
+        # groups being null.
+        on_off = next(g for g in reversed(comment.groups()) if g is not None)
+        if on_off == "off" and off_start is None:
+            off_start = comment.start()
+        elif on_off == "on" and off_start is not None:
+            off_ranges.append((off_start, comment.end()))
+            off_start = None
+    if off_start is not None:
+        off_ranges.append((off_start, len(src)))
+
+    def _off_range(start: int, end: int) -> bool:
+        return any(start >= rng[0] and end <= rng[1] for rng in off_ranges)
+
     @contextlib.contextmanager
     def _collect_error(match: Match[str]) -> Generator[None, None, None]:
         try:
@@ -110,13 +137,17 @@ def _collect_error(match: Match[str]) -> Generator[None, None, None]:
             errors.append(CodeBlockError(match.start(), e))
 
     def _md_match(match: Match[str]) -> str:
+        if _off_range(match.start(), match.end()):
+            return match.group(0)
         code = textwrap.dedent(match["code"])
         with _collect_error(match):
             code = black.format_str(code, mode=black_mode)
         code = textwrap.indent(code, match["indent"])
         return f'{match["before"]}{code}{match["after"]}'
 
     def _rst_match(match: Match[str]) -> str:
+        if _off_range(match.start(), match.end()):
+            return match.group(0)
         lang = match["lang"]
         if lang is not None and lang not in PYGMENTS_PY_LANGS:
             return match[0]
@@ -131,6 +162,8 @@ def _rst_match(match: Match[str]) -> str:
         return f'{match["before"]}{code.rstrip()}{trailing_ws}'
 
     def _rst_literal_blocks_match(match: Match[str]) -> str:
+        if _off_range(match.start(), match.end()):
+            return match.group(0)
         if not match["code"].strip():
             return match[0]
         min_indent = min(INDENT_RE.findall(match["code"]))
@@ -189,24 +222,32 @@ def finish_fragment() -> None:
         return code
 
     def _md_pycon_match(match: Match[str]) -> str:
+        if _off_range(match.start(), match.end()):
+            return match.group(0)
         code = _pycon_match(match)
         code = textwrap.indent(code, match["indent"])
         return f'{match["before"]}{code}{match["after"]}'
 
     def _rst_pycon_match(match: Match[str]) -> str:
+        if _off_range(match.start(), match.end()):
+            return match.group(0)
         code = _pycon_match(match)
         min_indent = min(INDENT_RE.findall(match["code"]))
         code = textwrap.indent(code, min_indent)
         return f'{match["before"]}{code}'
 
     def _latex_match(match: Match[str]) -> str:
+        if _off_range(match.start(), match.end()):
+            return match.group(0)
         code = textwrap.dedent(match["code"])
         with _collect_error(match):
             code = black.format_str(code, mode=black_mode)
         code = textwrap.indent(code, match["indent"])
         return f'{match["before"]}{code}{match["after"]}'
 
     def _latex_pycon_match(match: Match[str]) -> str:
+        if _off_range(match.start(), match.end()):
+            return match.group(0)
         code = _pycon_match(match)
         code = textwrap.indent(code, match["indent"])
         return f'{match["before"]}{code}{match["after"]}'

tests/test_blacken_docs.py

@@ -944,3 +944,150 @@ def test_format_src_rst_pycon_comment_before_promopt():
         "    # Comment about next line\n"
         "    >>> pass\n"
     )
+
+
+def test_on_off_comments_markdown_python():
+    before = (
+        "<!-- blacken-docs:off -->\n"
+        "```python\n"
+        "f(1,2,3)\n"
+        "```\n"
+        "<!-- blacken-docs:on -->\n"
+    )
+    after, _ = blacken_docs.format_str(before, BLACK_MODE)
+    assert after == before
+
+
+def test_on_off_comments_markdown_pycon():
+    before = (
+        "<!-- blacken-docs:off -->\n"
+        "```pycon\n"
+        ">>> f(1,2,3)\n"
+        "```\n"
+        "<!-- blacken-docs:on -->\n"
+    )
+    after, _ = blacken_docs.format_str(before, BLACK_MODE)
+    assert after == before
+
+
+def test_on_off_comments_rst_python():
+    before = (
+        "..\n"
+        "   blacken-docs:off\n"
+        ".. code-block:: python\n"
+        "\n"
+        "    f(1,2,3)\n"
+        "\n"
+        "..\n"
+        "   blacken-docs:on\n"
+    )
+    after, _ = blacken_docs.format_str(before, BLACK_MODE)
+    assert after == before
+
+
+def test_on_off_comments_rst_pycon():
+    before = (
+        "..\n"
+        "   blacken-docs:off\n"
+        ".. code-block:: pycon\n"
+        "\n"
+        "    >>> f(1,2,3)\n"
+        "\n"
+        "..\n"
+        "   blacken-docs:on\n"
+    )
+    after, _ = blacken_docs.format_str(before, BLACK_MODE)
+    assert after == before
+
+
+def test_on_off_comments_rst_literal():
+    before = (
+        "..\n"
+        "   blacken-docs:off\n"
+        "Example::\n"
+        "\n"
+        "    f(1,2,3)\n"
+        "\n"
+        "..\n"
+        "   blacken-docs:on\n"
+    )
+    after, _ = blacken_docs.format_str(before, BLACK_MODE, rst_literal_blocks=True)
+    assert after == before
+
+
+def test_on_off_comments_latex_python():
+    before = (
+        "% blacken-docs:off\n"
+        "\\begin{minted}{python}\n"
+        "f(1,2,3)\n"
+        "\\end{minted}\n"
+        "% blacken-docs:on\n"
+    )
+    after, _ = blacken_docs.format_str(before, BLACK_MODE)
+    assert after == before
+
+
+def test_on_off_comments_latex_pycon():
+    before = (
+        "% blacken-docs:off\n"
+        "\\begin{minted}{pycon}\n"
+        ">>> f(1,2,3)\n"
+        "\\end{minted}\n"
+        "% blacken-docs:on\n"
+    )
+    after, _ = blacken_docs.format_str(before, BLACK_MODE)
+    assert after == before
+
+
+def test_on_off_comments_flow():
+    before = (
+        "<!-- blacken-docs:on -->\n"  # ignored
+        "<!-- blacken-docs:off -->\n"
+        "<!-- blacken-docs:on -->\n"
+        "<!-- blacken-docs:on -->\n"  # ignored
+        "<!-- blacken-docs:off -->\n"
+        "<!-- blacken-docs:off -->\n"  # ignored
+        "```python\n"
+        "f(1,2,3)\n"
+        "```\n"  # no on comment, off until the end
+    )
+    after, _ = blacken_docs.format_str(before, BLACK_MODE)
+    assert after == before
+
+
+def test_on_off_comments_ranges():
+    before = (
+        "<!-- blacken-docs:off -->\n"
+        "```python\n"
+        "f(1,2,3)\n"
+        "```\n"
+        "<!-- blacken-docs:on -->\n"
+        "```python\n"
+        "f(1,2,3)\n"
+        "```\n"
+    )
+    after, _ = blacken_docs.format_str(before, BLACK_MODE)
+    assert after == (
+        "<!-- blacken-docs:off -->\n"
+        "```python\n"
+        "f(1,2,3)\n"
+        "```\n"
+        "<!-- blacken-docs:on -->\n"
+        "```python\n"
+        "f(1, 2, 3)\n"
+        "```\n"
+    )
+
+
+def test_on_off_comments_in_code_blocks():
+    before = (
+        "````md\n"
+        "<!-- blacken-docs:off -->\n"
+        "```python\n"
+        "f(1,2,3)\n"
+        "```\n"
+        "<!-- blacken-docs:on -->\n"
+        "````\n"
+    )
+    after, _ = blacken_docs.format_str(before, BLACK_MODE)
+    assert after == before