Separate base64-encoding from file handling

src/wireviz/Harness.py

@@ -18,7 +18,7 @@ from wireviz.wv_bom import manufacturer_info_field, component_table_entry, \
 from wireviz.wv_html import generate_html_output
 from wireviz.wv_helper import awg_equiv, mm2_equiv, tuplelist2tsv, flatten2d, \
     open_file_read, open_file_write
-from wireviz.svgembed import embed_svg_images
+from wireviz.svgembed import embed_svg_images_file
 
 
 class Harness:
@@ -358,7 +358,7 @@ class Harness:
             graph.render(filename=filename, view=view, cleanup=cleanup)
         graph.save(filename=f'{filename}.gv')
         if 'svg' in fmt:
-            embed_svg_images(f'{filename}.svg')
+            embed_svg_images_file(f'{filename}.svg')
         # bom output
         bomlist = bom_list(self.bom())
         with open_file_write(f'{filename}.bom.tsv') as file:

src/wireviz/svgembed.py

@@ -7,33 +7,42 @@ from pathlib import Path
 
 mime_subtype_replacements = {'jpg': 'jpeg', 'tif': 'tiff'}
 
-def embed_svg_images(filename_in: Path, overwrite: bool = True):
+def embed_svg_images(svg_in: str, base_path: Path):
+    # first, find any image references in SVG data, and cache the respective base64-encoded image
+    images_b64 = {}  # cache of base64-encoded images
+    re_xlink=re.compile(r"xlink:href=\"(?P<URL>.*?)\"", re.IGNORECASE)
+    for xlink in re_xlink.finditer(svg_in):
+        imgurl = xlink.group('URL')
+        if not imgurl in images_b64:  # only encode/cache every unique URL once
+            if not Path(imgurl).is_absolute():
+                imgurl_abs = (Path(base_path) / imgurl).resolve()
+            else:
+                imgurl_abs = Path(imgurl)
+
+            with open(imgurl_abs,'rb') as img:
+                img_bin = img.read()
+                img_b64 = base64.b64encode(img_bin)
+                img_str = img_b64.decode('utf-8')
+            images_b64[imgurl] = img_str
+    # second, replace links with the base64-encoded data
+    svg_out = svg_in
+    for url, b64 in images_b64.items():
+        svg_out = svg_out.replace(url, f'data:image/{get_mime_subtype(url)};base64, {b64}')
+
+    return svg_out
+
+
+def get_mime_subtype(filename: Path):
+    mime_subtype = str(Path(filename).suffix[1:]).lower()  # remove `.`
+    if mime_subtype in mime_subtype_replacements:
+        mime_subtype = mime_subtype_replacements[mime_subtype]
+    return mime_subtype
+
+
+def embed_svg_images_file(filename_in: Path, overwrite: bool = True):
     filename_in = Path(filename_in).resolve()
     filename_out = Path(f'{filename_in.with_suffix("")}.b64.svg')
-    images_b64 = {}  # cache base-64 encoded images
-    re_xlink=re.compile(r"xlink:href=\"(?P<URL>.*?)\"", re.IGNORECASE)
     with open(filename_in,'r') as file_in, open(filename_out,'w') as file_out:
-        for line in file_in:
+        file_out.write(embed_svg_images(file_in.read(), filename_in.parent))
-            for xlink in re_xlink.finditer(line):
-                imgurl = xlink.group('URL')
-                if not imgurl in images_b64:
-                    if not Path(imgurl).is_absolute():  # resolve relative image path
-                        imgurl_abs = (Path(filename_in).parent / imgurl).resolve()
-                    else:
-                        imgurl_abs = Path(imgurl)
-
-                    with open(imgurl_abs,'rb') as img:
-                        data_bin = img.read()
-                        data_b64 = base64.b64encode(data_bin)
-                        data_str = data_b64.decode('utf-8')
-                    images_b64[imgurl] = data_str
-
-                mime_subtype = str(imgurl_abs.suffix[1:]).lower()  # remove `.`
-                if mime_subtype in mime_subtype_replacements:
-                    mime_subtype = mime_subtype_replacements[mime_subtype]
-                line = line.replace(imgurl,
-                                    f'data:image/{mime_subtype};base64, {images_b64[imgurl]}')
-            file_out.write(line)
-
     if overwrite:
         filename_out.replace(filename_in)