Add support for on/off comments (#287)

Closes #193

---------

Co-authored-by: Adam Johnson <me@adamj.eu>

CHANGELOG.rst

@@ -2,6 +2,10 @@
 Changelog
 =========
 
+* Add support for on/off comments.
+
+  Thanks to Timothée Mazzucotelli in `PR #287 <https://github.com/adamchainz/blacken-docs/pull/287>`__.
+
 * Fix Markdown ``pycon`` formatting to allow formatting the rest of the file.
 
 1.17.0 (2024-06-29)

README.rst

@@ -155,6 +155,16 @@ And “pycon” blocks:
 
     ```
 
+Prevent formatting within a block using ``blacken-docs:off`` and ``blacken-docs:on`` comments:
+
+.. code-block:: markdown
+
+    <!-- blacken-docs:off -->
+    ```python
+    # whatever you want
+    ```
+    <!-- blacken-docs:on -->
+
 Within Python files, docstrings that contain Markdown code blocks may be reformatted:
 
 .. code-block:: python
@@ -189,6 +199,18 @@ In “pycon” blocks:
         ...     print("hello world")
         ...
 
+Prevent formatting within a block using ``blacken-docs:off`` and ``blacken-docs:on`` comments:
+
+.. code-block:: rst
+
+    .. blacken-docs:off
+
+    .. code-block:: python
+
+        # whatever you want
+
+    .. blacken-docs:on
+
 Use ``--rst-literal-blocks`` to also format `literal blocks <https://docutils.sourceforge.io/docs/ref/rst/restructuredtext.html#literal-blocks>`__:
 
 .. code-block:: rst
@@ -244,3 +266,13 @@ In PythonTeX blocks:
     def hello():
         print("hello world")
     \end{pycode}
+
+Prevent formatting within a block using ``blacken-docs:off`` and ``blacken-docs:on`` comments:
+
+.. code-block:: latex
+
+    % blacken-docs:off
+    \begin{minted}{python}
+    # whatever you want
+    \end{minted}
+    % blacken-docs:on

src/blacken_docs/__init__.py

@@ -4,6 +4,7 @@
 import contextlib
 import re
 import textwrap
+from bisect import bisect
 from typing import Generator
 from typing import Match
 from typing import Sequence
@@ -87,6 +88,16 @@
 )
 INDENT_RE = re.compile("^ +(?=[^ ])", re.MULTILINE)
 TRAILING_NL_RE = re.compile(r"\n+\Z", re.MULTILINE)
+ON_OFF = r"blacken-docs:(on|off)"
+ON_OFF_COMMENT_RE = re.compile(
+    # Markdown
+    rf"(?:^\s*<!-- {ON_OFF} -->$)|"
+    # rST
+    rf"(?:^\s*\.\. +{ON_OFF}$)|"
+    # LaTeX
+    rf"(?:^\s*% {ON_OFF}$)",
+    re.MULTILINE,
+)
 
 
 class CodeBlockError:
@@ -103,6 +114,29 @@ 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):
+        # Check for the "off" value across the multiple (on|off) groups.
+        if "off" in comment.groups():
+            if off_start is None:
+                off_start = comment.start()
+        else:
+            if 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 _within_off_range(code_range: tuple[int, int]) -> bool:
+        index = bisect(off_ranges, code_range)
+        try:
+            off_start, off_end = off_ranges[index - 1]
+        except IndexError:
+            return False
+        code_start, code_end = code_range
+        return code_start >= off_start and code_end <= off_end
+
     @contextlib.contextmanager
     def _collect_error(match: Match[str]) -> Generator[None, None, None]:
         try:
@@ -111,13 +145,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 _within_off_range(match.span()):
+            return match[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 _within_off_range(match.span()):
+            return match[0]
         lang = match["lang"]
         if lang is not None and lang not in PYGMENTS_PY_LANGS:
             return match[0]
@@ -132,6 +170,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 _within_off_range(match.span()):
+            return match[0]
         if not match["code"].strip():
             return match[0]
         min_indent = min(INDENT_RE.findall(match["code"]))
@@ -190,24 +230,32 @@ def finish_fragment() -> None:
         return code
 
     def _md_pycon_match(match: Match[str]) -> str:
+        if _within_off_range(match.span()):
+            return match[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 _within_off_range(match.span()):
+            return match[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 _within_off_range(match.span()):
+            return match[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 _within_off_range(match.span()):
+            return match[0]
         code = _pycon_match(match)
         code = textwrap.indent(code, match["indent"])
         return f'{match["before"]}{code}{match["after"]}'