openai/openai-python
Publicmirrored fromhttps://github.com/openai/openai-pythonAvailable
scripts/utils/ruffen-docs.py
167lines · modecode
| 1 | # fork of https://github.com/asottile/blacken-docs adapted for ruff |
| 2 | from __future__ import annotations |
| 3 | |
| 4 | import re |
| 5 | import sys |
| 6 | import argparse |
| 7 | import textwrap |
| 8 | import contextlib |
| 9 | import subprocess |
| 10 | from typing import Match, Optional, Sequence, Generator, NamedTuple, cast |
| 11 | |
| 12 | MD_RE = re.compile( |
| 13 | r"(?P<before>^(?P<indent> *)```\s*python\n)" r"(?P<code>.*?)" r"(?P<after>^(?P=indent)```\s*$)", |
| 14 | re.DOTALL | re.MULTILINE, |
| 15 | ) |
| 16 | MD_PYCON_RE = re.compile( |
| 17 | r"(?P<before>^(?P<indent> *)```\s*pycon\n)" r"(?P<code>.*?)" r"(?P<after>^(?P=indent)```.*$)", |
| 18 | re.DOTALL | re.MULTILINE, |
| 19 | ) |
| 20 | PYCON_PREFIX = ">>> " |
| 21 | PYCON_CONTINUATION_PREFIX = "..." |
| 22 | PYCON_CONTINUATION_RE = re.compile( |
| 23 | rf"^{re.escape(PYCON_CONTINUATION_PREFIX)}( |$)", |
| 24 | ) |
| 25 | DEFAULT_LINE_LENGTH = 100 |
| 26 | |
| 27 | |
| 28 | class CodeBlockError(NamedTuple): |
| 29 | offset: int |
| 30 | exc: Exception |
| 31 | |
| 32 | |
| 33 | def format_str( |
| 34 | src: str, |
| 35 | ) -> tuple[str, Sequence[CodeBlockError]]: |
| 36 | errors: list[CodeBlockError] = [] |
| 37 | |
| 38 | @contextlib.contextmanager |
| 39 | def _collect_error(match: Match[str]) -> Generator[None, None, None]: |
| 40 | try: |
| 41 | yield |
| 42 | except Exception as e: |
| 43 | errors.append(CodeBlockError(match.start(), e)) |
| 44 | |
| 45 | def _md_match(match: Match[str]) -> str: |
| 46 | code = textwrap.dedent(match["code"]) |
| 47 | with _collect_error(match): |
| 48 | code = format_code_block(code) |
| 49 | code = textwrap.indent(code, match["indent"]) |
| 50 | return f'{match["before"]}{code}{match["after"]}' |
| 51 | |
| 52 | def _pycon_match(match: Match[str]) -> str: |
| 53 | code = "" |
| 54 | fragment = cast(Optional[str], None) |
| 55 | |
| 56 | def finish_fragment() -> None: |
| 57 | nonlocal code |
| 58 | nonlocal fragment |
| 59 | |
| 60 | if fragment is not None: |
| 61 | with _collect_error(match): |
| 62 | fragment = format_code_block(fragment) |
| 63 | fragment_lines = fragment.splitlines() |
| 64 | code += f"{PYCON_PREFIX}{fragment_lines[0]}\n" |
| 65 | for line in fragment_lines[1:]: |
| 66 | # Skip blank lines to handle Black adding a blank above |
| 67 | # functions within blocks. A blank line would end the REPL |
| 68 | # continuation prompt. |
| 69 | # |
| 70 | # >>> if True: |
| 71 | # ... def f(): |
| 72 | # ... pass |
| 73 | # ... |
| 74 | if line: |
| 75 | code += f"{PYCON_CONTINUATION_PREFIX} {line}\n" |
| 76 | if fragment_lines[-1].startswith(" "): |
| 77 | code += f"{PYCON_CONTINUATION_PREFIX}\n" |
| 78 | fragment = None |
| 79 | |
| 80 | indentation = None |
| 81 | for line in match["code"].splitlines(): |
| 82 | orig_line, line = line, line.lstrip() |
| 83 | if indentation is None and line: |
| 84 | indentation = len(orig_line) - len(line) |
| 85 | continuation_match = PYCON_CONTINUATION_RE.match(line) |
| 86 | if continuation_match and fragment is not None: |
| 87 | fragment += line[continuation_match.end() :] + "\n" |
| 88 | else: |
| 89 | finish_fragment() |
| 90 | if line.startswith(PYCON_PREFIX): |
| 91 | fragment = line[len(PYCON_PREFIX) :] + "\n" |
| 92 | else: |
| 93 | code += orig_line[indentation:] + "\n" |
| 94 | finish_fragment() |
| 95 | return code |
| 96 | |
| 97 | def _md_pycon_match(match: Match[str]) -> str: |
| 98 | code = _pycon_match(match) |
| 99 | code = textwrap.indent(code, match["indent"]) |
| 100 | return f'{match["before"]}{code}{match["after"]}' |
| 101 | |
| 102 | src = MD_RE.sub(_md_match, src) |
| 103 | src = MD_PYCON_RE.sub(_md_pycon_match, src) |
| 104 | return src, errors |
| 105 | |
| 106 | |
| 107 | def format_code_block(code: str) -> str: |
| 108 | return subprocess.check_output( |
| 109 | [ |
| 110 | sys.executable, |
| 111 | "-m", |
| 112 | "ruff", |
| 113 | "format", |
| 114 | "--stdin-filename=script.py", |
| 115 | f"--line-length={DEFAULT_LINE_LENGTH}", |
| 116 | ], |
| 117 | encoding="utf-8", |
| 118 | input=code, |
| 119 | ) |
| 120 | |
| 121 | |
| 122 | def format_file( |
| 123 | filename: str, |
| 124 | skip_errors: bool, |
| 125 | ) -> int: |
| 126 | with open(filename, encoding="UTF-8") as f: |
| 127 | contents = f.read() |
| 128 | new_contents, errors = format_str(contents) |
| 129 | for error in errors: |
| 130 | lineno = contents[: error.offset].count("\n") + 1 |
| 131 | print(f"{filename}:{lineno}: code block parse error {error.exc}") |
| 132 | if errors and not skip_errors: |
| 133 | return 1 |
| 134 | if contents != new_contents: |
| 135 | print(f"{filename}: Rewriting...") |
| 136 | with open(filename, "w", encoding="UTF-8") as f: |
| 137 | f.write(new_contents) |
| 138 | return 0 |
| 139 | else: |
| 140 | return 0 |
| 141 | |
| 142 | |
| 143 | def main(argv: Sequence[str] | None = None) -> int: |
| 144 | parser = argparse.ArgumentParser() |
| 145 | parser.add_argument( |
| 146 | "-l", |
| 147 | "--line-length", |
| 148 | type=int, |
| 149 | default=DEFAULT_LINE_LENGTH, |
| 150 | ) |
| 151 | parser.add_argument( |
| 152 | "-S", |
| 153 | "--skip-string-normalization", |
| 154 | action="store_true", |
| 155 | ) |
| 156 | parser.add_argument("-E", "--skip-errors", action="store_true") |
| 157 | parser.add_argument("filenames", nargs="*") |
| 158 | args = parser.parse_args(argv) |
| 159 | |
| 160 | retv = 0 |
| 161 | for filename in args.filenames: |
| 162 | retv |= format_file(filename, skip_errors=args.skip_errors) |
| 163 | return retv |
| 164 | |
| 165 | |
| 166 | if __name__ == "__main__": |
| 167 | raise SystemExit(main()) |
| 168 | |