1#!/usr/bin/env python32# -*- coding: utf-8 -*-34import argparse5import os6from pathlib import Path7import sys8from typing import Any, Tuple910import yaml1112if __name__ == '__main__':13 sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))1415from wireviz import __version__16from wireviz.DataClasses import Metadata, Options, Tweak17from wireviz.Harness import Harness18from wireviz.wv_helper import expand, open_file_read192021def parse(yaml_input: str, file_out: (str, Path) = None, return_types: (None, str, Tuple[str]) = None) -> Any:22 """23 Parses yaml input string and does the high-level harness conversion2425 :param yaml_input: a string containing the yaml input data26 :param file_out:27 :param return_types: if None, then returns None; if the value is a string, then a28 corresponding data format will be returned; if the value is a tuple of strings,29 then for every valid format in the `return_types` tuple, another return type30 will be generated and returned in the same order; currently supports:31 - "png" - will return the PNG data32 - "svg" - will return the SVG data33 - "harness" - will return the `Harness` instance34 """3536 yaml_data = yaml.safe_load(yaml_input)3738 harness = Harness(39 metadata = Metadata(**yaml_data.get('metadata', {})),40 options = Options(**yaml_data.get('options', {})),41 tweak = Tweak(**yaml_data.get('tweak', {})),42 )4344 # When title is not given, either deduce it from filename, or use default text.45 if 'title' not in harness.metadata:46 if file_out is None:47 harness.metadata['title'] = "WireViz diagram and BOM"48 else:49 harness.metadata['title'] = Path(file_out).stem5051 # add items52 sections = ['connectors', 'cables', 'connections']53 types = [dict, dict, list]54 for sec, ty in zip(sections, types):55 if sec in yaml_data and type(yaml_data[sec]) == ty:56 if len(yaml_data[sec]) > 0:57 if ty == dict:58 for key, attribs in yaml_data[sec].items():59 # The Image dataclass might need to open an image file with a relative path.60 image = attribs.get('image')61 if isinstance(image, dict):62 image['gv_dir'] = Path(file_out if file_out else '').parent # Inject context6364 if sec == 'connectors':65 if not attribs.get('autogenerate', False):66 harness.add_connector(name=key, **attribs)67 elif sec == 'cables':68 harness.add_cable(name=key, **attribs)69 else:70 pass # section exists but is empty71 else: # section does not exist, create empty section72 if ty == dict:73 yaml_data[sec] = {}74 elif ty == list:75 yaml_data[sec] = []7677 # add connections7879 def check_designators(what, where): # helper function80 for i, x in enumerate(what):81 if x not in yaml_data[where[i]]:82 return False83 return True8485 autogenerated_ids = {}86 for connection in yaml_data['connections']:87 # find first component (potentially nested inside list or dict)88 first_item = connection[0]89 if isinstance(first_item, list):90 first_item = first_item[0]91 elif isinstance(first_item, dict):92 first_item = list(first_item.keys())[0]93 elif isinstance(first_item, str):94 pass9596 # check which section the first item belongs to97 alternating_sections = ['connectors','cables']98 for index, section in enumerate(alternating_sections):99 if first_item in yaml_data[section]:100 expected_index = index101 break102 else:103 raise Exception('First item not found anywhere.')104 expected_index = 1 - expected_index # flip once since it is flipped back at the *beginning* of every loop105106 # check that all iterable items (lists and dicts) are the same length107 # and that they are alternating between connectors and cables/bundles, starting with either108 itemcount = None109 for item in connection:110 expected_index = 1 - expected_index # make sure items alternate between connectors and cables111 expected_section = alternating_sections[expected_index]112 if isinstance(item, list):113 itemcount_new = len(item)114 for subitem in item:115 if not subitem in yaml_data[expected_section]:116 raise Exception(f'{subitem} is not in {expected_section}')117 elif isinstance(item, dict):118 if len(item.keys()) != 1:119 raise Exception('Dicts may contain only one key here!')120 itemcount_new = len(expand(list(item.values())[0]))121 subitem = list(item.keys())[0]122 if not subitem in yaml_data[expected_section]:123 raise Exception(f'{subitem} is not in {expected_section}')124 elif isinstance(item, str):125 if not item in yaml_data[expected_section]:126 raise Exception(f'{item} is not in {expected_section}')127 continue128 if itemcount is not None and itemcount_new != itemcount:129 raise Exception('All lists and dict lists must be the same length!')130 itemcount = itemcount_new131 if itemcount is None:132 raise Exception('No item revealed the number of connections to make!')133134 # populate connection list135 connection_list = []136 for i, item in enumerate(connection):137 if isinstance(item, str): # one single-pin component was specified138 sublist = []139 for i in range(1, itemcount + 1):140 if yaml_data['connectors'][item].get('autogenerate'):141 autogenerated_ids[item] = autogenerated_ids.get(item, 0) + 1142 new_id = f'_{item}_{autogenerated_ids[item]}'143 harness.add_connector(new_id, **yaml_data['connectors'][item])144 sublist.append([new_id, 1])145 else:146 sublist.append([item, 1])147 connection_list.append(sublist)148 elif isinstance(item, list): # a list of single-pin components were specified149 sublist = []150 for subitem in item:151 if yaml_data['connectors'][subitem].get('autogenerate'):152 autogenerated_ids[subitem] = autogenerated_ids.get(subitem, 0) + 1153 new_id = f'_{subitem}_{autogenerated_ids[subitem]}'154 harness.add_connector(new_id, **yaml_data['connectors'][subitem])155 sublist.append([new_id, 1])156 else:157 sublist.append([subitem, 1])158 connection_list.append(sublist)159 elif isinstance(item, dict): # a component with multiple pins was specified160 sublist = []161 id = list(item.keys())[0]162 pins = expand(list(item.values())[0])163 for pin in pins:164 sublist.append([id, pin])165 connection_list.append(sublist)166 else:167 raise Exception('Unexpected item in connection list')168169 # actually connect components using connection list170 for i, item in enumerate(connection_list):171 id = item[0][0] # TODO: make more elegant/robust/pythonic172 if id in harness.cables:173 for j, con in enumerate(item):174 if i == 0: # list started with a cable, no connector to join on left side175 from_name = None176 from_pin = None177 else:178 from_name = connection_list[i-1][j][0]179 from_pin = connection_list[i-1][j][1]180 via_name = item[j][0]181 via_pin = item[j][1]182 if i == len(connection_list) - 1: # list ends with a cable, no connector to join on right side183 to_name = None184 to_pin = None185 else:186 to_name = connection_list[i+1][j][0]187 to_pin = connection_list[i+1][j][1]188 harness.connect(from_name, from_pin, via_name, via_pin, to_name, to_pin)189190 if "additional_bom_items" in yaml_data:191 for line in yaml_data["additional_bom_items"]:192 harness.add_bom_item(line)193194 if file_out is not None:195 harness.output(filename=file_out, fmt=('png', 'svg'), view=False)196197 if return_types is not None:198 returns = []199 if isinstance(return_types, str): # only one return type specified200 return_types = [return_types]201202 return_types = [t.lower() for t in return_types]203204 for rt in return_types:205 if rt == 'png':206 returns.append(harness.png)207 if rt == 'svg':208 returns.append(harness.svg)209 if rt == 'harness':210 returns.append(harness)211212 return tuple(returns) if len(returns) != 1 else returns[0]213214215def parse_file(yaml_file: str, file_out: (str, Path) = None) -> None:216 with open_file_read(yaml_file) as file:217 yaml_input = file.read()218219 if not file_out:220 fn, fext = os.path.splitext(yaml_file)221 file_out = fn222 file_out = os.path.abspath(file_out)223224 parse(yaml_input, file_out=file_out)225226227def parse_cmdline():228 parser = argparse.ArgumentParser(229 description='Generate cable and wiring harness documentation from YAML descriptions',230 )231 parser.add_argument('-V', '--version', action='version', version='%(prog)s ' + __version__)232 parser.add_argument('-f', '--format', action='store', default='svg', type=str, help="Output formats.")233 parser.add_argument('input_file', action='store', type=str, metavar='YAML_FILE')234 parser.add_argument('-o', '--output_file', action='store', type=str, metavar='OUTPUT')235 # Not implemented: parser.add_argument('--generate-bom', action='store_true', default=True)236 parser.add_argument('--prepend-file', action='store', type=str, metavar='YAML_FILE')237 return parser.parse_args()238239240def main():241 args = parse_cmdline()242243 # input file244 input_file = args.input_file245 if input_file is None or input_file == '-':246 yaml_input = sys.stdin.read()247 else:248 if not os.path.exists(input_file):249 sys.stderr.write(f'Error: input file {input_file} inaccessible or does not exist, check path')250 sys.exit(1)251 with open_file_read(input_file) as fh:252 yaml_input = fh.read()253254 # prepend file255 prepend_file = args.prepend_file256 if prepend_file:257 if not os.path.exists(prepend_file):258 sys.stderr.write(f'Error: prepend input file {prepend_file} inaccessible or does not exist, check path')259 sys.exit(1)260 with open_file_read(prepend_file) as fh:261 prepend = fh.read()262 yaml_input = prepend + yaml_input263264 # output file265 output_file = args.output_file266 if output_file:267 if str(output_file) == "-":268 # write to stdout269 file_out = None270 else:271 file_out = os.path.abspath(output_file)272 else:273 if input_file == '-':274 # write to stdout275 file_out = None276 else:277 file_out = input_file278 pre, _ = os.path.splitext(file_out)279 file_out = os.path.abspath(pre) # extension will be added by graphviz output function280281 output = parse(yaml_input, file_out=file_out, return_types=args.format)282 if file_out is None:283 if isinstance(output, (bytes, bytearray)):284 sys.stdout.buffer.write(output)285 else:286 sys.stdout.write(output)287288289if __name__ == '__main__':290 main()