diff --git a/src/wireviz/DataClasses.py b/src/wireviz/DataClasses.py index d46dcb1d..56a38589 100644 --- a/src/wireviz/DataClasses.py +++ b/src/wireviz/DataClasses.py @@ -1,7 +1,7 @@ #!/usr/bin/env python3 # -*- coding: utf-8 -*- -from typing import Optional, List, Union +from typing import Optional, List, Union, Tuple from dataclasses import dataclass, field, InitVar from pathlib import Path from wireviz.wv_helper import int2tuple, aspect_ratio @@ -19,6 +19,8 @@ Color = str # Two-letter color name = Literal[wv_colors._color_hex.keys()] Colors = str # One or more two-letter color names (Color) concatenated into one string ColorScheme = str # Color scheme name = Literal[wv_colors.COLOR_CODES.keys()] +NoneOrMorePins = Union[Pin, Tuple[Pin, ...], None] # None, one, or a tuple of pins +OneOrMoreWires = Union[Wire, Tuple[Wire, ...]] # One or a tuple of wires @dataclass @@ -96,7 +98,7 @@ class Connector: ignore_in_bom: bool = False additional_components: List[AdditionalComponent] = field(default_factory=list) - def __post_init__(self): + def __post_init__(self) -> None: if isinstance(self.image, dict): self.image = Image(**self.image) @@ -148,7 +150,7 @@ def __post_init__(self): if isinstance(item, dict): self.additional_components[i] = AdditionalComponent(**item) - def activate_pin(self, pin): + def activate_pin(self, pin: Pin) -> None: self.visible_pins[pin] = True def get_qty_multiplier(self, qty_multiplier: Optional[ConnectorMultiplier]) -> int: @@ -186,7 +188,7 @@ class Cable: ignore_in_bom: bool = False additional_components: List[AdditionalComponent] = field(default_factory=list) - def __post_init__(self): + def __post_init__(self) -> None: if isinstance(self.image, dict): self.image = Image(**self.image) @@ -246,7 +248,9 @@ def __post_init__(self): if isinstance(item, dict): self.additional_components[i] = AdditionalComponent(**item) - def connect(self, from_name, from_pin, via_pin, to_name, to_pin): + # The *_pin arguments accept a tuple, but it seems not in use with the current code. + def connect(self, from_name: Optional[Designator], from_pin: NoneOrMorePins, via_pin: OneOrMoreWires, + to_name: Optional[Designator], to_pin: NoneOrMorePins) -> None: from_pin = int2tuple(from_pin) via_pin = int2tuple(via_pin) to_pin = int2tuple(to_pin)