onf/scripts/strip_rest_markup.py

"""Strip reST markup (x, :role:x, x) from docstrings and comments only.

Docstring spans come from the AST, comment spans from tokenize, so no code string literal is
touched. Prints a per-file count; --check re-parses and compares the docstring-stripped AST.
"""
from __future__ import annotations

import ast
import io
import re
import sys
import tokenize
from pathlib import Path

ROLE = re.compile(r":(?:class|meth|func|mod|attr|data|obj|exc|ref|term):`~?([^`]+)`")
DOUBLE = re.compile(r"``([^`]+)``")
SINGLE = re.compile(r"(?<!`)`([^`\n]+)`(?!`)")
AUTODOC = re.compile(r"(?m)^(\s*)#:( |$)")


def clean(text: str) -> str:
    text = ROLE.sub(r"\1", text)
    text = DOUBLE.sub(r"\1", text)
    text = AUTODOC.sub(r"\1#\2", text)
    return SINGLE.sub(r"\1", text)


def docstring_nodes(tree: ast.AST):
    for node in ast.walk(tree):
        if isinstance(node, (ast.Module, ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef)):
            body = getattr(node, "body", None)
            if body and isinstance(body[0], ast.Expr) and isinstance(body[0].value, ast.Constant) \
               and isinstance(body[0].value.value, str):
                yield body[0].value


def rewrite(src: str) -> tuple[str, int]:
    lines = src.splitlines(keepends=True)
    offsets = [0]
    for line in lines:
        offsets.append(offsets[-1] + len(line))

    def pos(row: int, col: int) -> int:
        return offsets[row - 1] + col

    spans: list[tuple[int, int]] = []
    for node in docstring_nodes(ast.parse(src)):
        spans.append((pos(node.lineno, node.col_offset), pos(node.end_lineno, node.end_col_offset)))
    for tok in tokenize.generate_tokens(io.StringIO(src).readline):
        if tok.type == tokenize.COMMENT:
            spans.append((pos(tok.start[0], tok.start[1]), pos(tok.end[0], tok.end[1])))

    spans.sort()
    out, prev, n = [], 0, 0
    for lo, hi in spans:
        if lo < prev:
            continue
        original = src[lo:hi]
        replaced = clean(original)
        n += original != replaced
        out.append(src[prev:lo])
        out.append(replaced)
        prev = hi
    out.append(src[prev:])
    return "".join(out), n


def stripped_ast(src: str) -> str:
    tree = ast.parse(src)
    for node in ast.walk(tree):
        if isinstance(node, (ast.Module, ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef)):
            body = getattr(node, "body", None)
            if body and isinstance(body[0], ast.Expr) and isinstance(body[0].value, ast.Constant) \
               and isinstance(body[0].value.value, str):
                node.body = body[1:] or [ast.Pass()]
    return ast.dump(tree)


def main(roots: list[str]) -> int:
    changed = failed = 0
    for root in roots:
        for path in sorted(Path(root).rglob("*.py")):
            if "__pycache__" in path.parts:
                continue
            src = path.read_text()
            new, n = rewrite(src)
            if not n:
                continue
            if stripped_ast(src) != stripped_ast(new):
                print(f"AST CHANGED, refusing: {path}")
                failed += 1
                continue
            path.write_text(new)
            changed += 1
            print(f"{n:4d} spans  {path}")
    print(f"\n{changed} files rewritten, {failed} refused")
    return 1 if failed else 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))