# SPDX-FileCopyrightText: 2022-present deepset GmbH # # SPDX-License-Identifier: Apache-2.0 """ Pre-commit hook that runs ruff format on Python code blocks in Markdown/MDX files. Uses the ruff configuration from pyproject.toml automatically. """ import argparse import logging import os import re import subprocess import sys import tempfile logger = logging.getLogger(__name__) logger.setLevel(logging.INFO) handler = logging.StreamHandler(sys.stderr) handler.setFormatter(logging.Formatter("%(message)s")) logger.addHandler(handler) # ANSI color codes (disabled when stderr is not a terminal) _USE_COLOR = hasattr(sys.stderr, "isatty") and sys.stderr.isatty() PYTHON_FENCE_RE = re.compile( r"(?P^```python\s*\n)" r"(?P.*?)" r"(?P^```\s*$)", re.MULTILINE | re.DOTALL, ) def _color(code: str, text: str) -> str: """Colorize the text""" if _USE_COLOR: return f"\033[{code}m{text}\033[0m" return text def _find_tool(name: str) -> str: """Find a tool installed in the same virtualenv as the running Python.""" return os.path.join(os.path.dirname(sys.executable), name) def _ruff(code: str, *, line_length: int) -> str: return subprocess.run( [ _find_tool("ruff"), "format", f"--line-length={line_length}", "--config", "format.skip-magic-trailing-comma = false", "--stdin-filename", "block.py", "-", ], input=code, capture_output=True, text=True, check=True, ).stdout def _add_trailing_commas(code: str) -> str: """Add trailing commas to multi-line expressions using add-trailing-comma.""" with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f: f.write(code) tmpfile = f.name try: subprocess.run([_find_tool("add-trailing-comma"), tmpfile], capture_output=True, check=False) with open(tmpfile) as f: return f.read() finally: os.unlink(tmpfile) def _format_code_block(match: re.Match, *, line_length: int, path: str) -> str: """Format a single code block""" code = match.group("code") try: # 1. ruff format (may create new multi-line expressions) # 2. add trailing commas to all multi-line expressions # 3. ruff format again (respects trailing commas, ensures stable output) formatted = _ruff(_add_trailing_commas(_ruff(code, line_length=line_length)), line_length=line_length) except subprocess.CalledProcessError as exc: snippet = code.strip().splitlines() preview = "\n".join(snippet[:5]) if len(snippet) > 5: preview += f"\n... ({len(snippet) - 5} more lines)" logger.warning( "%s %s\n%s\n%s %s", _color("33", "WARNING:"), _color("1", path), _color("2", preview), _color("31", "ruff stderr:"), exc.stderr.strip(), ) return match.group(0) return match.group("before") + formatted + match.group("after") def main() -> int: """Main entrypoint""" parser = argparse.ArgumentParser() parser.add_argument("--line-length", type=int, default=120) parser.add_argument("files", nargs="*") args = parser.parse_args() ret = 0 for path in args.files: with open(path) as f: original = f.read() new = PYTHON_FENCE_RE.sub( lambda m, path=path: _format_code_block(m, line_length=args.line_length, path=path), original ) if new != original: with open(path, "w") as f: f.write(new) logger.debug("%s %s", _color("32", "Rewriting:"), _color("1", path)) ret = 1 return ret if __name__ == "__main__": raise SystemExit(main())