diff --git a/src/wireviz/Harness.py b/src/wireviz/Harness.py index d7e0859..4d9fd51 100644 --- a/src/wireviz/Harness.py +++ b/src/wireviz/Harness.py @@ -328,23 +328,38 @@ class Harness: @property def csv(self): - raise NotImplementedError + return bom_helper.generate_bom_outputs( + self.bom_list(), + 'csv' + ) @property def csv_excel(self): - raise NotImplementedError + return bom_helper.generate_bom_outputs( + self.bom_list(), + 'csv_excel' + ) @property def csv_unix(self): - raise NotImplementedError + return bom_helper.generate_bom_outputs( + self.bom_list(), + 'csv_unix' + ) @property def tsv(self): - raise NotImplementedError + return bom_helper.generate_bom_outputs( + self.bom_list(), + 'tsv' + ) @property def tsv_excel(self): - raise NotImplementedError + return bom_helper.generate_bom_outputs( + self.bom_list(), + 'tsv_excel' + ) def output(self, filename: (str, Path), view: bool = False, cleanup: bool = True, fmt: tuple = ('pdf', )) -> None: # graphical output @@ -356,7 +371,7 @@ class Harness: # bom output bom_list = self.bom_list() # todo: support user choices of BOM format (probably also graphviz outputs, html outputs) - bom_helper.generate_bom_outputs(filename,bom_list, [bom_helper.WIREVIZ_TSV, bom_helper.EXCEL_CSV]) + bom_helper.generate_bom_outputs(filename, bom_list, [bom_helper.WIREVIZ_TSV, bom_helper.EXCEL_CSV]) # HTML output with open_file_write(f'{filename}.html') as file: file.write('\n') diff --git a/src/wireviz/bom_helper.py b/src/wireviz/bom_helper.py index 2bcc27c..c6b3cbc 100644 --- a/src/wireviz/bom_helper.py +++ b/src/wireviz/bom_helper.py @@ -2,12 +2,10 @@ # -*- coding: utf-8 -*- import csv -from wireviz import wv_helper -from wireviz.wv_helper import open_file_write +from io import StringIO + +from wireviz import wv_helper -EXCEL_CSV = csv.excel -EXCEL_TSV = csv.excel_tab -UNIX_CSV = csv.unix_dialect WIREVIZ_TSV = type('Wireviz BOM', (csv.Dialect, object), dict( delimiter='\t', doublequote=True, @@ -20,31 +18,26 @@ WIREVIZ_TSV = type('Wireviz BOM', (csv.Dialect, object), dict( )) csv.register_dialect('Wireviz BOM', WIREVIZ_TSV) -_csv_formats = { EXCEL_CSV, UNIX_CSV } -_tsv_formats = { EXCEL_TSV, WIREVIZ_TSV } -_csv_ext = '.bom.csv' -_tsv_ext = '.bom.tsv' +def generate_bom_outputs(bomdata, dialect='tsv'): + dialect = dialect.strip().lower() + dialect_lookup = { + 'csv': csv.unix_dialect, + 'csv_unix': csv.unix_dialect, + 'csv_excel': csv.excel, + 'tsv': WIREVIZ_TSV, + 'tsv_excel': csv.excel_tab + } + valid_dialects = [k for k in dialect_lookup.keys()] + if dialect not in valid_dialects: + raise ValueError(f'dialect "{dialect}" not supported') + output = StringIO() + writer = csv.writer(output, dialect=dialect_lookup[dialect]) + writer.writerows(wv_helper.flatten2d(bomdata)) -def generate_bom_outputs(base_filename, bomdata, formats=None): - if formats is None: - formats = [EXCEL_CSV, WIREVIZ_TSV] - elif isinstance(formats, csv.Dialect): - formats = [formats] - elif not isinstance(formats, list): - raise TypeError - expanded_csv_names = len(_csv_formats.intersection(set(formats))) > 1 - expanded_tsv_names = len(_tsv_formats.intersection(set(formats))) > 1 + output.seek(0) - for fmt in formats: - if fmt in _csv_formats: - file = csv.writer(open_file_write(base_filename + ("_" + fmt.__name__ if expanded_csv_names else "") + _csv_ext, fmt.lineterminator), fmt) - - elif fmt in _tsv_formats: - file = csv.writer(open_file_write(base_filename + ("_"+fmt.__name__ if expanded_tsv_names else "") + _tsv_ext, fmt.lineterminator), fmt) - else: - raise KeyError("Unknown BOM Format Specified") - file.writerows(wv_helper.flatten2d(bomdata)) + return bytes(output.read(), encoding='utf-8') # TODO: Possibly refactor other BOM output operations, such as HTML, into here? \ No newline at end of file diff --git a/src/wireviz/wireviz.py b/src/wireviz/wireviz.py index 0797ef4..9f9697b 100755 --- a/src/wireviz/wireviz.py +++ b/src/wireviz/wireviz.py @@ -249,7 +249,13 @@ def main(input_file, prepend, out): filedatas = parse(yaml_input, return_types=out) if isinstance(filedatas, tuple): for ext, data in zip(out, filedatas): - fname = f'{base_file_name}.{ext}' + if 'csv' in ext: + extension = 'bom.csv' + elif 'tsv' in ext: + extension = 'bom.tsv' + else: + extension = ext + fname = f'{base_file_name}.{extension}' with open(fname, 'wb') as f: f.write(data) else: