Coverage for diagrams / code_diagrams / _source_marker.py: 100.00%
104 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-09-14 22:07 +0000
1"""Idempotently insert ``# Flowchart:`` markers into source files.
3The marker sits right under the module docstring and points readers to
4the generated flowchart artifacts. Re-running the generator is safe:
5existing marker blocks (identified by a sentinel) are replaced in place
6rather than duplicated.
7"""
9from __future__ import annotations
11import ast
12import re
13import subprocess
14import sys
15import warnings
16from collections import defaultdict
17from pathlib import Path
19from diagrams.code_diagrams._renderer import RenderedTarget
21SENTINEL = "pyflowchart-code-diagram"
22_BLOCK_RE = re.compile(
23 rf"(?s)# <{re.escape(SENTINEL)}> BEGIN.*?# <{re.escape(SENTINEL)}> END\n?",
24)
27def upsert_markers(
28 results: list[RenderedTarget],
29 *,
30 project_root: Path,
31) -> None:
32 """Add or refresh a pointer comment in every source file we charted.
34 Multiple targets from the same source file collapse into a single
35 comment block that lists every flowchart produced for that file.
36 After writing the marker, each touched file is normalised with
37 ``ruff format`` so the resulting layout is formatter-stable —
38 otherwise the marker's leading/trailing blank lines can compose
39 with the source file's existing PEP 8 spacing into three-blank-line
40 runs that break ``ruff format --check``.
41 """
42 by_source: dict[Path, list[RenderedTarget]] = defaultdict(list)
43 for result in results:
44 by_source[project_root / result.target.source].append(result)
46 touched: list[Path] = []
47 for source_path, source_results in by_source.items():
48 if _update_file(
49 source_path=source_path,
50 results=source_results,
51 project_root=project_root,
52 ):
53 touched.append(source_path)
55 if touched:
56 _ruff_format(touched, project_root=project_root)
59def _ruff_format(paths: list[Path], *, project_root: Path) -> None:
60 """Run ``ruff format`` on ``paths`` so the marker insertion is formatter-stable.
62 Uses ``python -m ruff`` from the current interpreter so the check
63 works whether ``ruff`` is on PATH or only importable. Silently
64 no-ops if ``ruff`` isn't importable at all — the generator still
65 works, the contributor just has to run ``ruff format`` themselves later.
66 """
67 try:
68 import ruff # noqa: F401
69 except ImportError:
70 warnings.warn(
71 "ruff is not installed — skipping post-marker normalisation. "
72 "Install with ``pip install -e '.[diagrams]'``; the marker block "
73 "may still land in a shape ruff later reformats.",
74 stacklevel=2,
75 )
76 return
78 rels = [str(p.relative_to(project_root)) for p in paths]
79 # Invoke ruff directly so we inherit its exit code + stdout.
80 subprocess.run( # noqa: S603 — args are fully-known paths we just generated
81 [sys.executable, "-m", "ruff", "format", "--quiet", *rels],
82 cwd=str(project_root),
83 check=True,
84 )
87def _update_file(
88 *,
89 source_path: Path,
90 results: list[RenderedTarget],
91 project_root: Path,
92) -> bool:
93 """Insert or replace the marker block in ``source_path``.
95 Returns ``True`` iff the file was actually modified. Implementation
96 note: we always *strip* any existing marker first, then re-insert
97 at the current ``_insertion_point`` offset. Doing an in-place
98 ``re.sub`` when the block is present is correct for the idempotent
99 case, but fails silently when the placement rules change (e.g.
100 when we moved the block from "after ``from __future__ import ...``"
101 to "after all imports"). A strip-then-insert pipeline also makes
102 ``--skip-marker=False`` + an upstream schema change land the marker
103 in the right spot without the user having to run a separate
104 cleanup pass.
106 We don't try to normalise whitespace here — :func:`upsert_markers`
107 runs ``ruff format`` on every touched file after all insertions
108 complete, which handles the PEP 8 blank-line spacing consistently.
109 """
110 original = source_path.read_text(encoding="utf-8")
111 stripped = strip_markers_from(original)
112 block = _format_block(results=results, project_root=project_root)
114 insertion_point = _insertion_point(stripped)
115 updated = stripped[:insertion_point] + block + stripped[insertion_point:]
117 if updated != original:
118 source_path.write_text(updated, encoding="utf-8")
119 print(f" 🖋 marker inserted/refreshed in {source_path.relative_to(project_root)}")
120 return True
121 return False
124def strip_markers_from(source: str) -> str:
125 """Return ``source`` with any existing marker block removed.
127 Kept as a public helper so the CLI ``--strip-markers`` flag can
128 reuse the exact same regex. If no marker block is present the
129 source is returned unchanged. When a marker block is removed we
130 collapse the resulting *four*-or-more consecutive newlines down
131 to a single three-newline run (i.e. two blank lines) — that
132 preserves the PEP 8 ``two-blank-lines-between-top-level-defs``
133 requirement ruff format enforces, which would otherwise be broken
134 when the marker block lived between two top-level defs and removing
135 it fused their trailing and leading blank-line padding into a
136 four-newline run. Files with legitimate triple-blank-line runs
137 unrelated to a marker are left unchanged.
138 """
139 if SENTINEL not in source:
140 return source
141 without_block = _BLOCK_RE.sub("", source)
142 return re.sub(r"\n{4,}", "\n\n\n", without_block)
145def strip_all_markers(project_root: Path) -> int:
146 """Remove every marker block under ``project_root``.
148 Walks the standard source roots — ``app.py``, ``cli/``, ``gco/``,
149 ``gco_mcp/``, and ``lambda/`` (excluding the kubectl-applier-simple-build
150 and helm-installer-build packaged bundles) — and rewrites any file
151 that actually contains a marker. Files without the sentinel are
152 left untouched (even if they have triple-blank-line runs
153 unrelated to this feature). Returns the number of files modified.
154 """
155 modified = 0
156 search_roots: list[Path] = [
157 project_root / "app.py",
158 *(project_root / "cli").rglob("*.py"),
159 *(project_root / "gco").rglob("*.py"),
160 *(project_root / "gco_mcp").rglob("*.py"),
161 *(project_root / "lambda").rglob("*.py"),
162 ]
163 skip_fragments = ("kubectl-applier-simple-build", "helm-installer-build")
164 for source_path in search_roots:
165 if not source_path.is_file():
166 continue
167 if any(frag in str(source_path) for frag in skip_fragments):
168 continue
169 original = source_path.read_text(encoding="utf-8")
170 if SENTINEL not in original:
171 continue
172 stripped = strip_markers_from(original)
173 if stripped != original:
174 source_path.write_text(stripped, encoding="utf-8")
175 print(f" 🧹 stripped marker from {source_path.relative_to(project_root)}")
176 modified += 1
177 return modified
180def _format_block(
181 *,
182 results: list[RenderedTarget],
183 project_root: Path,
184) -> str:
185 """Build the comment block that points at the generated artifacts.
187 The block is preceded *and* followed by a blank line so it cleanly
188 separates from the surrounding statements — otherwise ruff's
189 ``I001`` rule treats the comment as part of the import block above,
190 and PEP 8 enforcement complains about the single blank line between
191 the marker and a class/def below. Two blank-line separators work in
192 every context we insert into (after docstring, after ``from __future__``
193 imports, after a regular import block, before a class/def/module-level
194 statement).
195 """
196 lines = ["", f"# <{SENTINEL}> BEGIN - auto-inserted, do not edit"]
197 source_commits = {result.source_commit for result in results}
198 if len(source_commits) != 1:
199 raise ValueError("source marker results must share one Git source commit")
200 source_commit = next(iter(source_commits))
201 lines.append(f"# Generated at (UTC): {results[0].generated_at}")
202 lines.append(f"# Generated from Git commit: {source_commit}")
203 lines.append("# Flowchart(s) generated from this file:")
204 for result in results:
205 html_rel = result.html_path.relative_to(project_root)
206 lines.append(f"# * ``{result.target.function}`` -> ``{html_rel}``")
207 if result.png_path is not None:
208 png_rel = result.png_path.relative_to(project_root)
209 lines.append(f"# (PNG: ``{png_rel}``)")
210 lines.append(
211 "# Regenerate with ``SOURCE_DATE_EPOCH=<unix-seconds> "
212 "GCO_DIAGRAM_SOURCE_COMMIT=<40-char-sha> "
213 "python diagrams/generate.py --code-only``.",
214 )
215 lines.append(f"# <{SENTINEL}> END")
216 # Trailing "" plus the final "\n" from ``join`` ensures the block
217 # ends with an empty line — combined with whatever line-terminator
218 # is already present in the source at the insertion point, this
219 # gives us the two-blank-line separator ruff format expects before
220 # the next class or def.
221 lines.append("")
222 return "\n".join(lines) + "\n"
225def _insertion_point(source: str) -> int:
226 """Return the character offset where the marker block should go.
228 Places the block immediately after the module docstring and the
229 full block of top-level imports (``import``, ``from`` — including
230 ``from __future__ import …``), but before the first real
231 statement. This placement keeps ruff's import sorter happy: it
232 groups consecutive imports, and a comment block slotted in the
233 middle of that group would be treated as a section boundary that
234 forces reordering.
236 Falls back to offset 0 if the file has no docstring and no imports.
237 """
238 try:
239 tree = ast.parse(source)
240 except SyntaxError: # pragma: no cover - defensive
241 return 0
243 last_prelude_end_line = 0
245 # Module docstring: the first statement is a bare string expression.
246 body_iter = iter(tree.body)
247 first = next(body_iter, None)
248 if (
249 first is not None
250 and isinstance(first, ast.Expr)
251 and isinstance(first.value, ast.Constant)
252 and isinstance(first.value.value, str)
253 ):
254 last_prelude_end_line = first.end_lineno or 0
255 else:
256 # No docstring — the first statement (if any) is already
257 # imports / code, so reset the iterator to include it.
258 body_iter = iter(tree.body)
260 # Walk every subsequent top-level ``import`` / ``from ... import ...``
261 # statement. The first non-import node terminates the prelude.
262 for node in body_iter:
263 if isinstance(node, ast.Import | ast.ImportFrom):
264 last_prelude_end_line = max(last_prelude_end_line, node.end_lineno or 0)
265 else:
266 break
268 if last_prelude_end_line == 0:
269 return 0
271 # Convert line number (1-indexed, inclusive) -> char offset after
272 # the newline that terminates that line.
273 offset = 0
274 for _ in range(last_prelude_end_line):
275 newline = source.find("\n", offset)
276 if newline == -1:
277 return len(source)
278 offset = newline + 1
279 return offset