Add save/load from png

Author: Jacob Erickson

src/wireviz/wireviz.py

@@ -8,6 +8,8 @@
 from typing import Any, Tuple
 
 import yaml
+from PIL import Image
+from PIL.PngImagePlugin import PngInfo
 
 if __name__ == '__main__':
     sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
@@ -181,6 +183,7 @@ def check_designators(what, where): # helper function
 
     if file_out is not None:
         harness.output(filename=file_out, fmt=('png', 'svg'), view=False)
+        save_yaml_to_png(file_out,yaml_input)
 
     if return_types is not None:
         returns = []
@@ -223,6 +226,16 @@ def parse_cmdline():
     parser.add_argument('--prepend-file', action='store', type=str, metavar='YAML_FILE')
     return parser.parse_args()
 
+def save_yaml_to_png(file_out,yaml_input):
+    with Image.open(fp=f'{file_out}.png') as im:
+        txt = PngInfo()
+        txt.add_itxt('yaml',yaml_input,zip=True)
+        im.save(fp=f'{file_out}.png',pnginfo=txt)
+
+def read_yaml_from_png(file_in):
+    with Image.open(fp=f'{file_in}.png') as im:
+        im.load()
+        return im.text['yaml']
 
 def main():
 
@@ -232,16 +245,19 @@ def main():
         print(f'Error: input file {args.input_file} inaccessible or does not exist, check path')
         sys.exit(1)
 
-    with open_file_read(args.input_file) as fh:
-        yaml_input = fh.read()
-
-    if args.prepend_file:
-        if not os.path.exists(args.prepend_file):
-            print(f'Error: prepend input file {args.prepend_file} inaccessible or does not exist, check path')
-            sys.exit(1)
-        with open_file_read(args.prepend_file) as fh:
-            prepend = fh.read()
-            yaml_input = prepend + yaml_input
+    if ".png" in args.input_file:
+        yaml_input = read_yaml_from_png(args.input_file.replace('.png',''))
+    else:
+        with open_file_read(args.input_file) as fh:
+            yaml_input = fh.read()
+
+        if args.prepend_file:
+            if not os.path.exists(args.prepend_file):
+                print(f'Error: prepend input file {args.prepend_file} inaccessible or does not exist, check path')
+                sys.exit(1)
+            with open_file_read(args.prepend_file) as fh:
+                prepend = fh.read()
+                yaml_input = prepend + yaml_input
 
     if not args.output_file:
         file_out = args.input_file