From c84e6fb3adfb42f90ebb759eff43d01d3a2f1ec8 Mon Sep 17 00:00:00 2001
From: Ryan Malloy <ryan@supported.systems>
Date: Fri, 13 Feb 2026 00:27:09 -0700
Subject: [PATCH] Embed YAML source in PNG metadata for round-trip capability
 (port of upstream PR #234)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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.
---
 src/wireviz/wireviz.py         | 11 ++++--
 src/wireviz/wv_png_metadata.py | 63 ++++++++++++++++++++++++++++++++++
 2 files changed, 71 insertions(+), 3 deletions(-)
 create mode 100644 src/wireviz/wv_png_metadata.py

diff --git a/src/wireviz/wireviz.py b/src/wireviz/wireviz.py
index fac305e..5706ad6 100755
--- a/src/wireviz/wireviz.py
+++ b/src/wireviz/wireviz.py
@@ -14,6 +14,7 @@ if __name__ == "__main__":
 
 from wireviz.wv_dataclasses import AUTOGENERATED_PREFIX, Metadata, Options, Tweak
 from wireviz.wv_harness import Harness
+from wireviz.wv_png_metadata import save_yaml_to_png
 from wireviz.wv_utils import (
     expand,
     file_read_text,
@@ -88,7 +89,7 @@ def parse(
     if not output_formats and not return_types:
         raise Exception("No output formats or return types specified")
 
-    yaml_data, yaml_file = _get_yaml_data_and_path(inp)
+    yaml_data, yaml_file, yaml_str = _get_yaml_data_and_path(inp)
     if not isinstance(yaml_data, dict):
         raise TypeError(
             f"Expected a dict as top-level YAML input, but got: {type(yaml_data)}"
@@ -394,6 +395,9 @@ def parse(
 
     if output_formats:
         harness.output(filename=output_file, fmt=output_formats, view=False)
+        # embed YAML source into PNG for round-trip capability
+        if "png" in output_formats and yaml_str:
+            save_yaml_to_png(output_file, yaml_str)
 
     if return_types:
         returns = []
@@ -413,7 +417,7 @@ def parse(
         return tuple(returns) if len(returns) != 1 else returns[0]
 
 
-def _get_yaml_data_and_path(inp: Union[str, Path, Dict]) -> Tuple[Dict, Path]:
+def _get_yaml_data_and_path(inp: Union[str, Path, Dict]) -> Tuple[Dict, Path, str]:
     # determine whether inp is a file path, a YAML string, or a Dict
     if not isinstance(inp, Dict):  # received a str or a Path
         try:
@@ -441,7 +445,8 @@ def _get_yaml_data_and_path(inp: Union[str, Path, Dict]) -> Tuple[Dict, Path]:
         # received a Dict, use as-is
         yaml_data = inp
         yaml_path = None
-    return yaml_data, yaml_path
+        yaml_str = None
+    return yaml_data, yaml_path, yaml_str
 
 
 def _get_output_dir(input_file: Path, default_output_dir: Path) -> Path:
diff --git a/src/wireviz/wv_png_metadata.py b/src/wireviz/wv_png_metadata.py
new file mode 100644
index 0000000..f4a0445
--- /dev/null
+++ b/src/wireviz/wv_png_metadata.py
@@ -0,0 +1,63 @@
+# -*- 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