Ryan Malloy
c84e6fb3ad
PNGs now contain the original YAML harness definition as compressed iTXT metadata. Use read_yaml_from_png() to extract it — share a PNG, recipient can regenerate or edit the harness.
64 lines
1.8 KiB
Python
64 lines
1.8 KiB
Python
# -*- coding: utf-8 -*-
|
|
|
|
"""Embed and extract YAML source data in PNG metadata (iTXT chunks)."""
|
|
|
|
from pathlib import Path
|
|
from typing import Optional, Tuple
|
|
|
|
from PIL import Image
|
|
from PIL.PngImagePlugin import PngInfo
|
|
|
|
|
|
PNG_KEY_YAML = "wireviz_yaml"
|
|
PNG_KEY_PREPEND = "wireviz_prepend_yaml"
|
|
|
|
|
|
def save_yaml_to_png(
|
|
png_path: Path,
|
|
yaml_input: str,
|
|
prepend_input: str = "",
|
|
) -> None:
|
|
"""Save YAML source as compressed iTXT metadata in a PNG file."""
|
|
png_path = Path(png_path)
|
|
if not png_path.suffix == ".png":
|
|
png_path = png_path.with_suffix(".png")
|
|
if not png_path.exists():
|
|
return
|
|
|
|
with Image.open(fp=png_path) as im:
|
|
txt = PngInfo()
|
|
txt.add_itxt(PNG_KEY_YAML, yaml_input, zip=True)
|
|
if prepend_input:
|
|
txt.add_itxt(PNG_KEY_PREPEND, prepend_input, zip=True)
|
|
im.save(fp=png_path, pnginfo=txt)
|
|
|
|
|
|
def read_yaml_from_png(png_path: Path) -> Tuple[str, Optional[str]]:
|
|
"""Extract YAML source from a PNG file's iTXT metadata.
|
|
|
|
Returns (yaml_input, prepend_input) where prepend_input may be None.
|
|
"""
|
|
png_path = Path(png_path)
|
|
if not png_path.suffix == ".png":
|
|
png_path = png_path.with_suffix(".png")
|
|
|
|
with Image.open(fp=png_path) as im:
|
|
im.load()
|
|
yaml_input = im.text.get(PNG_KEY_YAML, "")
|
|
prepend_input = im.text.get(PNG_KEY_PREPEND)
|
|
|
|
return yaml_input, prepend_input
|
|
|
|
|
|
def has_yaml_metadata(png_path: Path) -> bool:
|
|
"""Check if a PNG file contains embedded WireViz YAML data."""
|
|
png_path = Path(png_path)
|
|
if not png_path.suffix == ".png":
|
|
png_path = png_path.with_suffix(".png")
|
|
if not png_path.exists():
|
|
return False
|
|
|
|
with Image.open(fp=png_path) as im:
|
|
im.load()
|
|
return PNG_KEY_YAML in im.text
|