From 6f9bb67d023a99f3d860594969da6790d95f8f90 Mon Sep 17 00:00:00 2001
From: Daniel Rojas <github@danielrojas.net>
Date: Sat, 16 Oct 2021 22:36:42 +0200
Subject: [PATCH] Embed images into SVG output

---
 src/wireviz/Harness.py  | 23 +++++++++++-------
 src/wireviz/svgembed.py | 52 +++++++++++++++++++++++++++++++++++++++++
 src/wireviz/wv_html.py  |  2 +-
 3 files changed, 67 insertions(+), 10 deletions(-)
 create mode 100644 src/wireviz/svgembed.py

diff --git a/src/wireviz/Harness.py b/src/wireviz/Harness.py
index 5acda02c..de7163a2 100644
--- a/src/wireviz/Harness.py
+++ b/src/wireviz/Harness.py
@@ -20,6 +20,7 @@
     Tweak,
     Side,
 )
+from wireviz.svgembed import embed_svg_images_file
 from wireviz.wv_bom import (
     HEADER_MPN,
     HEADER_PN,
@@ -646,13 +647,9 @@ def png(self):
 
     @property
     def svg(self):
-        from io import BytesIO
-
         graph = self.graph
-        data = BytesIO()
-        data.write(graph.pipe(format="svg"))
-        data.seek(0)
-        return data.read()
+        return embed_svg_images(graph.pipe(format="svg").decode("utf-8"), Path.cwd())
+
 
     def output(
         self,
@@ -671,9 +668,14 @@ def output(
             if f in ("png", "svg", "html"):
                 if f == "html":  # if HTML format is specified,
                     f = "svg"  # generate SVG for embedding into HTML
+                # SVG file will be renamed/deleted later
+                _filename = f"{filename}.tmp" if f == "svg" else filename
                 # TODO: prevent rendering SVG twice when both SVG and HTML are specified
                 graph.format = f
-                graph.render(filename=filename, view=view, cleanup=cleanup)
+                graph.render(filename=_filename, view=view, cleanup=cleanup)
+        # embed images into SVG output
+        if "svg" in fmt or "html" in fmt:
+            embed_svg_images_file(f"{filename}.tmp.svg")
         # GraphViz output
         if "gv" in fmt:
             graph.save(filename=f"{filename}.gv")
@@ -692,8 +694,11 @@ def output(
             # TODO: implement PDF output
             print("PDF output is not yet supported")
         # delete SVG if not needed
-        if "html" in fmt and not "svg" in fmt and not svg_already_exists:
-            Path(f"{filename}.svg").unlink()
+        if "html" in fmt and not "svg" in fmt:
+            # SVG file was just needed to generate HTML
+            Path(f"{filename}.tmp.svg").unlink()
+        elif "svg" in fmt:
+            Path(f"{filename}.tmp.svg").replace(f"{filename}.svg")
 
     def bom(self):
         if not self._bom:
diff --git a/src/wireviz/svgembed.py b/src/wireviz/svgembed.py
new file mode 100644
index 00000000..ab6b9f1e
--- /dev/null
+++ b/src/wireviz/svgembed.py
@@ -0,0 +1,52 @@
+# -*- coding: utf-8 -*-
+
+import base64
+import re
+from pathlib import Path
+from typing import Union
+
+mime_subtype_replacements = {"jpg": "jpeg", "tif": "tiff"}
+
+
+def embed_svg_images(svg_in: str, base_path: Union[str, Path] = Path.cwd()) -> str:
+    images_b64 = {}  # cache of base64-encoded images
+
+    def image_tag(pre: str, url: str, post: str) -> str:
+        return f'<image{pre} xlink:href="{url}"{post}>'
+
+    def replace(match: re.Match) -> str:
+        imgurl = match["URL"]
+        if not imgurl in images_b64:  # only encode/cache every unique URL once
+            imgurl_abs = (Path(base_path) / imgurl).resolve()
+            image = imgurl_abs.read_bytes()
+            images_b64[imgurl] = base64.b64encode(image).decode("utf-8")
+        return image_tag(
+            match["PRE"] or "",
+            f"data:image/{get_mime_subtype(imgurl)};base64, {images_b64[imgurl]}",
+            match["POST"] or "",
+        )
+
+    pattern = re.compile(
+        image_tag(r"(?P<PRE> [^>]*?)?", r'(?P<URL>[^"]*?)', r"(?P<POST> [^>]*?)?"),
+        re.IGNORECASE,
+    )
+    return pattern.sub(replace, svg_in)
+
+
+def get_mime_subtype(filename: Union[str, Path]) -> str:
+    mime_subtype = Path(filename).suffix.lstrip(".").lower()
+    if mime_subtype in mime_subtype_replacements:
+        mime_subtype = mime_subtype_replacements[mime_subtype]
+    return mime_subtype
+
+
+def embed_svg_images_file(
+    filename_in: Union[str, Path], overwrite: bool = True
+) -> None:
+    filename_in = Path(filename_in).resolve()
+    filename_out = filename_in.with_suffix(".b64.svg")
+    filename_out.write_text(
+        embed_svg_images(filename_in.read_text(), filename_in.parent)
+    )
+    if overwrite:
+        filename_out.replace(filename_in)
diff --git a/src/wireviz/wv_html.py b/src/wireviz/wv_html.py
index b30bdeeb..15342668 100644
--- a/src/wireviz/wv_html.py
+++ b/src/wireviz/wv_html.py
@@ -37,7 +37,7 @@ def generate_html_output(
     html = open_file_read(templatefile).read()
 
     # embed SVG diagram
-    with open_file_read(f"{filename}.svg") as file:
+    with open_file_read(f"{filename}.tmp.svg") as file:
         svgdata = re.sub(
             "^<[?]xml [^?>]*[?]>[^<]*<!DOCTYPE [^>]*>",
             "<!-- XML and DOCTYPE declarations from SVG file removed -->",