diff --git a/src/wireviz/Harness.py b/src/wireviz/Harness.py index ca9f8e6..95e7a2a 100644 --- a/src/wireviz/Harness.py +++ b/src/wireviz/Harness.py @@ -20,16 +20,16 @@ class Harness: self.cables = {} self.additional_bom_items = [] - def add_connector(self, name, *args, **kwargs): + def add_connector(self, name: str, *args, **kwargs) -> None: self.connectors[name] = Connector(name, *args, **kwargs) - def add_cable(self, name, *args, **kwargs): + def add_cable(self, name: str, *args, **kwargs) -> None: self.cables[name] = Cable(name, *args, **kwargs) - def add_bom_item(self, item): + def add_bom_item(self, item: dict) -> None: self.additional_bom_items.append(item) - def connect(self, from_name, from_pin, via_name, via_pin, to_name, to_pin): + def connect(self, from_name: str, from_pin: (int, str), via_name: str, via_pin: (int, str), to_name: str, to_pin: (int, str)) -> None: for (name, pin) in zip([from_name, to_name], [from_pin, to_pin]): # check from and to connectors if name is not None and name in self.connectors: connector = self.connectors[name] @@ -58,7 +58,7 @@ class Harness: if to_name in self.connectors: self.connectors[to_name].activate_pin(to_pin) - def create_graph(self): + def create_graph(self) -> Graph: dot = Graph() dot.body.append('// Graph generated by WireViz') dot.body.append('// https://github.com/formatc1702/WireViz') @@ -280,7 +280,7 @@ class Harness: data.seek(0) return data.read() - def output(self, filename: (str, Path), view=False, cleanup=True, fmt=('pdf', )): + def output(self, filename: (str, Path), view: bool = False, cleanup: bool = True, fmt: tuple = ('pdf', )) -> None: # graphical output graph = self.create_graph() for f in fmt: diff --git a/src/wireviz/wireviz.py b/src/wireviz/wireviz.py index d52d3b1..6c9a186 100755 --- a/src/wireviz/wireviz.py +++ b/src/wireviz/wireviz.py @@ -3,8 +3,9 @@ import argparse import os +from pathlib import Path import sys -from typing import Tuple +from typing import Any, Tuple import yaml @@ -16,7 +17,7 @@ from wireviz.Harness import Harness from wireviz.wv_helper import expand -def parse(yaml_input, file_out=None, return_types: (None, str, Tuple[str]) = None): +def parse(yaml_input: str, file_out: (str, Path) = None, return_types: (None, str, Tuple[str]) = None) -> Any: """ Parses yaml input string and does the high-level harness conversion @@ -196,7 +197,7 @@ def parse(yaml_input, file_out=None, return_types: (None, str, Tuple[str]) = Non return tuple(returns) if len(returns) != 1 else returns[0] -def parse_file(yaml_file, file_out=None): +def parse_file(yaml_file: str, file_out: (str, Path) = None) -> None: with open(yaml_file, 'r') as file: yaml_input = file.read()