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