Coverage for diagrams/code_diagrams/

1"""pyflowchart + Playwright rendering helpers.

3Splitting the rendering concerns out of

4:mod:`diagrams.code_diagrams.generate` keeps the entry point small and

5makes it easy to unit-test the path math without importing Playwright.

6"""

8from __future__ import annotations

10import contextlib

11import hashlib

12import math

13import sys

14import warnings

15from dataclasses import dataclass

16from pathlib import Path

18from diagrams.code_diagrams._targets import Target

21@dataclass(frozen=True)

22class RenderedTarget:

23 """Output paths produced for a single :class:`Target`.

25 Paths are all absolute so callers don't need to know where the

26 project root lives.

27 """

29 target: Target

30 html_path: Path

31 png_path: Path | None

32 """``None`` if PNG rendering was skipped or failed."""

33 generated_at: str

34 """Invocation-wide ISO-8601 UTC generation timestamp."""

35 source_commit: str

36 """Exact Git commit containing the marker-stripped charted source."""

39def render_all(

40 *,

41 targets: list[Target],

42 project_root: Path,

43 output_dir: Path,

44 render_png: bool,

45 generated_at: str,

46 source_commit: str,

47) -> list[RenderedTarget]:

48 """Render every target, returning where each output landed."""

49 _require_pyflowchart()

50 renderer = _make_png_renderer() if render_png else None

51 try:

52 results: list[RenderedTarget] = []

53 for target in targets:

54 result = _render_one(

55 target=target,

56 project_root=project_root,

57 output_dir=output_dir,

58 renderer=renderer,

59 generated_at=generated_at,

60 source_commit=source_commit,

61 )

62 results.append(result)

63 return results

64 finally:

65 if renderer is not None:

66 renderer.close()

69def _render_one(

70 *,

71 target: Target,

72 project_root: Path,

73 output_dir: Path,

74 renderer: _PlaywrightRenderer | None,

75 generated_at: str,

76 source_commit: str,

77) -> RenderedTarget:

78 """Render a single target and return its output paths."""

79 from pyflowchart import Flowchart, output_html # local import: optional dep

81 source_path = (project_root / target.source).resolve()

82 source = source_path.read_text(encoding="utf-8")

83 print(f"\n🧭 {target.source}::{target.function}")

85 # ``Flowchart.from_code`` handles simplification and the field selector

86 # in one call; ``inner=True`` gives a control-flow chart of the body.

87 flowchart = Flowchart.from_code(source, field=target.function, inner=target.inner)

88 dsl = flowchart.flowchart()

90 stem = _output_stem_for(target, output_dir=output_dir)

91 html_path = stem.parent / f"{stem.name}.html"

92 png_path = stem.parent / f"{stem.name}.png"

93 html_path.parent.mkdir(parents=True, exist_ok=True)

95 title = target.title or f"{target.source}::{target.function}"

96 output_html(str(html_path), title, dsl)

97 # pyflowchart's HTML template includes trailing spaces on some generated

98 # lines. Normalize every artifact here so regeneration remains compatible

99 # with ``git diff --check`` across pyflowchart releases.

100 html = _annotate_generated_html(

101 html_path.read_text(encoding="utf-8"),

102 generated_at=generated_at,

103 source_commit=source_commit,

104 )

105 html_path.write_text(

106 "\n".join(line.rstrip() for line in html.splitlines()) + "\n",

107 encoding="utf-8",

108 )

109 print(f" ✓ HTML {html_path.relative_to(project_root)}")

110

111 # Never retain a PNG generated under an older invocation timestamp. Delete

112 # it before either attempting a fresh render or returning HTML-only output;

113 # a failed/skip-PNG run must not leave a mixed-age artifact set behind.

114 if png_path.is_file():

115 png_path.unlink()

116 print(f" 🧹 PNG removed stale {png_path.relative_to(project_root)}")

117

118 if renderer is not None:

119 ok = renderer.render(html_path=html_path, png_path=png_path)

120 if ok:

121 print(f" ✓ PNG {png_path.relative_to(project_root)}")

122 return RenderedTarget(

123 target=target,

124 html_path=html_path,

125 png_path=png_path,

126 generated_at=generated_at,

127 source_commit=source_commit,

128 )

129 return RenderedTarget(

130 target=target,

131 html_path=html_path,

132 png_path=None,

133 generated_at=generated_at,

134 source_commit=source_commit,

135 )

136

137 # ``--skip-png`` or Playwright unavailable. The stale artifact was removed

138 # above, so README/source markers accurately describe this as HTML-only.

139 return RenderedTarget(

140 target=target,

141 html_path=html_path,

142 png_path=None,

143 generated_at=generated_at,

144 source_commit=source_commit,

145 )

146

147

148def _annotate_generated_html(html: str, *, generated_at: str, source_commit: str) -> str:

149 """Add visible timestamp and deterministic flow-content metadata.

150

151 The visible wrapper intentionally contains the flowchart canvas so the

152 Playwright screenshot includes both the timestamp and a digest of the

153 pre-annotation pyflowchart HTML. Even when flowchart.js collapses changed

154 source into an otherwise identical SVG node, the paired PNG visibly changes

155 with the source-derived digest. The digest is a freshness marker, not a

156 cross-platform PNG byte-reproducibility claim.

157 """

158 charset = ' <meta charset="utf-8">'

159 canvas = ' <div id="canvas"></div>'

160 if charset not in html or canvas not in html:

161 raise RuntimeError("pyflowchart HTML template no longer matches the annotator")

162

163 flow_digest = hashlib.sha256(html.encode("utf-8")).hexdigest()[:16]

164 meta = "\n".join(

165 [

166 f' <meta name="gco-generated-at" content="{generated_at}">',

167 f' <meta name="gco-source-commit" content="{source_commit}">',

168 f' <meta name="gco-flow-digest" content="{flow_digest}">',

169 ]

170 )

171 artifact = "\n".join(

172 [

173 f" ",

174 f" ",

175 ' <div id="generated-artifact" style="display: inline-block; padding: 12px; background: #fff;">',

176 ' <p style="margin: 0 0 6px; color: #444; font: 14px Helvetica, sans-serif;">',

177 f' Generated at (UTC): <time datetime="{generated_at}">{generated_at}</time>',

178 " </p>",

179 ' <p style="margin: 0 0 6px; color: #555; font: 12px Helvetica, sans-serif;">',

180 f" Source commit: <code>{source_commit}</code>",

181 " </p>",

182 ' <p style="margin: 0 0 10px; color: #666; font: 12px Helvetica, sans-serif;">',

183 f" Flow content SHA-256: <code>{flow_digest}</code>",

184 " </p>",

185 ' <div id="canvas"></div>',

186 " </div>",

187 ],

188 )

189 return html.replace(charset, f"{charset}\n{meta}", 1).replace(canvas, artifact, 1)

190

191

192def _output_stem_for(target: Target, *, output_dir: Path) -> Path:

193 """Compute the output path stem for ``target`` (no suffix).

194

195 The output mirrors the source layout so large trees stay navigable.

196 For a source at ``lambda/analytics-presigned-url/handler.py`` with

197 function ``lambda_handler``, the stem is

198 ``<output_dir>/lambda/analytics-presigned-url/handler.lambda_handler``

199 (callers add ``.html`` / ``.png`` themselves; we cannot use

200 :meth:`Path.with_suffix` here because ``.lambda_handler`` would be

201 interpreted as a suffix and stripped).

202 """

203 src = Path(target.source)

204 return output_dir / src.parent / f"{src.stem}.{target.slug()}"

205

206

207def prune_orphaned_artifacts(*, targets: list[Target], output_dir: Path) -> list[Path]:

208 """Delete generated HTML/PNG files that no longer have a target.

209

210 Only full-catalog runs call this helper. Restricting cleanup to the two

211 generated suffixes preserves the generator source, README, and unrelated

212 files while removing renamed-source trees and retired targets. Empty

213 directories left behind by those artifacts are removed bottom-up.

214 """

215 expected: set[Path] = set()

216 for target in targets:

217 stem = _output_stem_for(target, output_dir=output_dir)

218 expected.update({stem.parent / f"{stem.name}.html", stem.parent / f"{stem.name}.png"})

219

220 removed: list[Path] = []

221 for artifact in sorted(

222 path

223 for path in output_dir.rglob("*")

224 if path.is_file() and path.suffix in {".html", ".png"}

225 ):

226 if artifact not in expected:

227 artifact.unlink()

228 removed.append(artifact)

229 print(f" 🧹 removed obsolete artifact {artifact.relative_to(output_dir)}")

230

231 directories = sorted(

232 (path for path in output_dir.rglob("*") if path.is_dir()),

233 key=lambda path: len(path.parts),

234 reverse=True,

235 )

236 for directory in directories:

237 if directory.name == "__pycache__":

238 continue

239 with contextlib.suppress(OSError):

240 directory.rmdir()

241

242 return removed

243

244

245def write_readme(results: list[RenderedTarget], *, output_dir: Path) -> None:

246 """(Re)generate ``code_diagrams/README.md`` with a grouped index."""

247 from diagrams.code_diagrams._readme import render_readme

248

249 readme_path = output_dir / "README.md"

250 content = render_readme(results, output_dir=output_dir)

251 readme_path.write_text(content, encoding="utf-8")

252 print(f"\n📝 Wrote {readme_path}")

253

254

255def _require_pyflowchart() -> None:

256 try:

257 import pyflowchart # noqa: F401

258 except ImportError as exc:

259 sys.exit(

260 "pyflowchart is not installed. Install the project's "

261 "``diagrams`` extra: ``pip install -e '.[diagrams]'``. "

262 f"(underlying error: {exc})"

263 )

264

265

266def _screenshot_scale(width: float, height: float) -> float:

267 """Bound Chromium screenshots by both dimension and total pixel area."""

268 max_css_dimension = 8_000

269 max_css_area = 20_000_000

270 return min(

271 1.0,

272 max_css_dimension / max(width, height),

273 math.sqrt(max_css_area / (width * height)),

274 )

275

276

277class _PlaywrightRenderer:

278 """Thin wrapper that keeps a single Playwright browser alive.

279

280 We intentionally open/close the browser at the batch boundary (not

281 per-target) so rendering dozens of targets doesn't pay the ~1s

282 browser start-up cost each time.

283 """

284

285 def __init__(self) -> None: # pragma: no cover - requires browser

286 from playwright.sync_api import sync_playwright

287

288 self._pw = sync_playwright().start()

289 self._browser = self._pw.chromium.launch()

290

291 def render(self, *, html_path: Path, png_path: Path) -> bool:

292 """Screenshot the flowchart SVG from ``html_path`` into ``png_path``.

293

294 Returns ``True`` on success.

295 """

296 # ``Error`` is Playwright's base exception; ``TimeoutError`` subclasses

297 # it. We catch both so a single un-renderable diagram degrades to

298 # HTML-only (``png_path=None``) instead of aborting the whole batch.

299 from playwright.sync_api import Error as PlaywrightError # pragma: no cover

300 from playwright.sync_api import TimeoutError as PwTimeout # pragma: no cover

301

302 page = self._browser.new_page(

303 viewport={"width": 2400, "height": 1800},

304 device_scale_factor=2,

305 )

306 try:

307 page.goto(html_path.absolute().as_uri())

308 # flowchart.js renders into ``<div id="canvas">`` — wait for

309 # the first child SVG node before screenshotting. Otherwise

310 # we capture the empty pre-render container.

311 page.wait_for_function(

312 "document.querySelector('#canvas svg') !== null",

313 timeout=30_000,

314 )

315 page.wait_for_timeout(500) # give layout a beat to settle

316 locator = page.locator("#canvas svg")

317 box = locator.bounding_box()

318 # Chromium rejects screenshots above roughly 32k physical pixels

319 # on either axis. A CSS transform changes only painting, not the

320 # element bounds Playwright passes to captureScreenshot, so resize

321 # the SVG viewport itself. The viewBox preserves all chart content.

322 # The area cap also avoids allocating several hundred megapixels

323 # for unusually wide-and-tall control-flow charts.

324 factor = _screenshot_scale(box["width"], box["height"]) if box is not None else 1.0

325 if box is not None and factor < 1.0:

326 locator.evaluate(

327 """(svg, size) => {

328 if (!svg.hasAttribute('viewBox')) {

329 svg.setAttribute(

330 'viewBox',

331 `0 0 ${size.sourceWidth} ${size.sourceHeight}`,

332 );

333 }

334 svg.setAttribute('width', size.width);

335 svg.setAttribute('height', size.height);

336 svg.style.width = `${size.width}px`;

337 svg.style.height = `${size.height}px`;

338 }""",

339 {

340 "sourceWidth": box["width"],

341 "sourceHeight": box["height"],

342 "width": box["width"] * factor,

343 "height": box["height"] * factor,

344 },

345 )

346 page.wait_for_timeout(100)

347 page.locator("#generated-artifact").screenshot(path=str(png_path))

348 return True

349 except PwTimeout as exc:

350 warnings.warn(

351 f"Playwright timed out rendering {html_path}: {exc}",

352 stacklevel=2,

353 )

354 return False

355 except PlaywrightError as exc:

356 # Most common cause: the flowchart is taller/wider than

357 # Chromium's maximum screenshot dimensions (~32k px), so

358 # ``Page.captureScreenshot`` returns "Unable to capture

359 # screenshot". The interactive HTML is still written and

360 # remains the primary artifact for these large diagrams.

361 warnings.warn(

362 f"Playwright could not screenshot {html_path} "

363 f"(diagram may exceed Chromium's max size): {exc}",

364 stacklevel=2,

365 )

366 return False

367 finally:

368 page.close()

369

370 def close(self) -> None:

371 """Shut down the browser and Playwright driver."""

372 try:

373 self._browser.close()

374 finally:

375 self._pw.stop()

376

377

378def _make_png_renderer() -> _PlaywrightRenderer | None:

379 """Best-effort Playwright initialisation.

380

381 Returns ``None`` (with a warning) if Playwright or its browsers

382 aren't installed, so the generator still produces the interactive

383 HTML even in environments that can't run Chromium.

384 """

385 try:

386 return _PlaywrightRenderer()

387 except ImportError:

388 warnings.warn(

389 "Playwright not installed — skipping PNG rendering. "

390 "Install with ``pip install -e '.[diagrams]'`` and then "

391 "``playwright install chromium``.",

392 stacklevel=2,

393 )

394 return None

395 except Exception as exc: # pragma: no cover - environment dependent

396 warnings.warn(

397 f"Playwright failed to start ({exc}) — skipping PNG rendering. "

398 "Run ``playwright install chromium`` to fetch the browser.",

399 stacklevel=2,

400 )

401 return None

Coverage for diagrams / code_diagrams / _renderer.py: 100.00%

133 statements