diff --git a/src/wireviz/Harness.py b/src/wireviz/Harness.py index 03336dd..eba63e2 100644 --- a/src/wireviz/Harness.py +++ b/src/wireviz/Harness.py @@ -255,6 +255,15 @@ class Harness: return dot + @property + def png(self): + from io import BytesIO + graph = self.create_graph() + data = BytesIO() + data.write(graph.pipe(format='png')) + data.seek(0) + return data.read() + def output(self, filename, directory='_output', view=False, cleanup=True, fmt='pdf', gen_bom=False): # graphical output graph = self.create_graph() diff --git a/src/wireviz/wireviz.py b/src/wireviz/wireviz.py index 27d9d55..75dfd0c 100755 --- a/src/wireviz/wireviz.py +++ b/src/wireviz/wireviz.py @@ -4,6 +4,7 @@ import argparse import os import sys +from typing import Tuple import yaml @@ -14,7 +15,19 @@ if __name__ == '__main__': from wireviz.Harness import Harness -def parse(yaml_input, file_out=None, generate_bom=False): +def parse(yaml_input, file_out=None, generate_bom=False, return_types: (None, str, Tuple[str]) = None): + """ + Parses yaml input string and does the high-level harness conversion + + :param yaml_input: a string containing the yaml input data + :param file_out: + :param generate_bom: + :param return_types: if None, then returns None; if the value is a string, then a + corresponding data format will be returned; if the value is a tuple of strings, + then for every valid format in the `return_types` tuple, another return type + will be generated and returned in the same order; currently supports only 'png', + but could easily support other types + """ yaml_data = yaml.safe_load(yaml_input) @@ -180,7 +193,21 @@ def parse(yaml_input, file_out=None, generate_bom=False): else: raise Exception('Wrong number of connection parameters') - harness.output(filename=file_out, fmt=('png', 'svg'), gen_bom=generate_bom, view=False) + if file_out is not None: + harness.output(filename=file_out, fmt=('png', 'svg'), gen_bom=generate_bom, view=False) + + if return_types is not None: + # gather the harness data and return it as a single data type + if isinstance(return_types, str): + if return_types.lower() == 'png': + return harness.png + + else: + # gather the harness data and return it as a series of tuples + returns = [] + if 'png' in return_types: + returns.append(harness.png) + return tuple(returns) def parse_file(yaml_file, file_out=None, generate_bom=False):