generate-wireviz-template/generate_wireviz_template.py

#! /usr/bin/env python3

"""
generate_wireviz_template.py: Generate a wireviz template for the connections
in the given multiboard assembly.
"""

import argparse
import io
import logging
import os
import re
import yaml

from dataclasses import dataclass, field
from typing import Mapping, Sequence, TextIO, Optional

from allspice import AllSpice, Repository
from allspice.utils.list_components import list_components, ComponentAttributes

ATTRIBUTES_TO_CHECK = [
    "_reference",
    "_logical_reference",
    "Designator",
    "Part Reference",
    "Reference",
]
"""
Attributes to check for in the component attributes to determine if the
component is a connector.
"""

CONNECTOR_DESIGNATOR_REGEX = r"^(J\d+|P\d+|CN\d+|X\w\d+)$"
"""
If a designator starts with:

- J: Jack
- P: Plug
- CN: Connector
- X(A-Z): Socket connector. The second character has to be present and be
  immediately followed by a digit distinguish it form X(TAL).

Source: https://en.wikipedia.org/wiki/Reference_designator

In each case, we accept designators of the pattern (letter)(numbers). 
We do not need to consider multi-part components here as we request combined
components from py-allspice. 
"""

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)


@dataclass
class MultiboardAssembly:
    @dataclass
    class Dependency:
        repo: str
        """Repository URL"""

        source_path: str
        reference: Optional[str] = None
        variant: Optional[str] = None
        metadata: Optional[Mapping[str, str]] = None
        quantity: int = 1

        @property
        def owner(self) -> str:
            return self.repo.split("/")[-2]

        @property
        def name(self) -> str:
            return self.repo.split("/")[-1]

    Variants = Mapping[str, Sequence[Dependency]]

    name: str
    description: str
    dependencies: Sequence[Dependency] = field(default_factory=list)
    variants: Variants = field(default_factory=dict)

    @staticmethod
    def load_yaml(file: TextIO) -> "MultiboardAssembly":
        data = yaml.safe_load(file)

        dependencies = [
            MultiboardAssembly.Dependency(**dep) for dep in data.get("dependencies", [])
        ]

        variants = {}
        for variant_name, variant_deps in data.get("variants", {}).items():
            variant_dependencies = [
                MultiboardAssembly.Dependency(**dep) for dep in variant_deps
            ]
            variants[variant_name] = variant_dependencies

        return MultiboardAssembly(
            name=data["name"],
            description=data["description"],
            dependencies=dependencies,
            variants=variants,
        )


@dataclass
class WirevizConnector:
    designator: str
    notes: str
    pin_labels: Sequence[str]

    def to_dict(self) -> dict:
        """
        Returns a dict that can be used to generate a wireviz template.
        """

        return {
            self.designator: {
                "notes": self.notes,
                "pinlabels": list(self.pin_labels),
            }
        }


def get_dependency_components(
    client: AllSpice,
    dependency: MultiboardAssembly.Dependency,
) -> list[ComponentAttributes]:
    """
    Get the components list of a dependency.

    Right now, this assumes the dependency is on the same AllSpice Hub instance
    as the client.
    """

    try:
        repo_name, repo_owner = dependency.repo.split("/")[-1:-3:-1]
    except ValueError:
        raise ValueError(f"Invalid repository URL: {dependency.repo}")

    repo = Repository.request(client, repo_owner, repo_name)
    return list_components(
        client,
        repo,
        dependency.source_path,
        variant=dependency.variant,
        ref=dependency.reference or repo.default_branch,
        combine_multi_part=True,
    )


def get_component_designator(component: ComponentAttributes) -> str:
    for attribute in ATTRIBUTES_TO_CHECK:
        attr_value = component.get(attribute)
        if attr_value is None or not isinstance(attr_value, str):
            continue
        else:
            return attr_value
    return ""


def find_connector_components(
    components: list[ComponentAttributes],
) -> list[ComponentAttributes]:
    connector_components = []

    for component in components:
        for attribute in ATTRIBUTES_TO_CHECK:
            attr_value = component.get(attribute)
            if attr_value is None or not isinstance(attr_value, str):
                continue

            if re.match(CONNECTOR_DESIGNATOR_REGEX, attr_value):
                connector_components.append(component)
                break

    return connector_components


def convert_component_to_wireviz_connector(
    component: ComponentAttributes,
    source_dependency: MultiboardAssembly.Dependency,
) -> list[WirevizConnector]:
    component_designator = get_component_designator(component)
    combined_designator = f"{source_dependency.name}_{component_designator}"
    # For Altium components, we get a better result indexing by part_id instead
    # of name, so we'll use fallbacks.
    component_name = component.get("_part_id") or component.get("_name") or ""

    notes = f"{component_name} from {source_dependency.repo}"
    if source_dependency.variant:
        notes += f" (variant {source_dependency.variant})"

    # The ComponentAttributes type doesn't know that pins is a list
    pins: list[dict] = component.get("_pins", [])  # type: ignore
    logger.debug("Found %d pins for component %s", len(pins), component_designator)
    pin_labels = []
    if pins:
        for pin in pins:
            pin_name = pin.get("name") or pin.get("designator")
            if pin_name:
                pin_labels.append(pin_name)
            else:
                logger.warn("Pin %s does not have a name or designator; skipping.", pin)

    if source_dependency.quantity > 1:
        return [
            WirevizConnector(
                designator=f"{combined_designator}_{i}",
                notes=notes + f" (part {i})",
                pin_labels=pin_labels,
            )
            for i in range(1, source_dependency.quantity + 1)
        ]
    else:
        return [
            WirevizConnector(
                designator=combined_designator,
                notes=notes,
                pin_labels=pin_labels,
            )
        ]


def main():
    arg_parser = argparse.ArgumentParser(description=__doc__)
    arg_parser.add_argument(
        "--assembly-file",
        required=True,
        help="Path to the assembly file to generate the wireviz template for.",
    )
    arg_parser.add_argument(
        "--variant",
        required=False,
        help="Variant to use for the assembly file. If blank or not given, the default variant is used.",
    )
    arg_parser.add_argument(
        "--output-file",
        required=True,
        help="Path to the output file to write the wireviz template to.",
    )
    arg_parser.add_argument(
        "--repository",
        required=True,
        help="Repository name in owner/repo formmat.",
    )
    arg_parser.add_argument(
        "--ref",
        required=True,
        help="Reference to use for the repository.",
    )
    arg_parser.add_argument(
        "--allspice-hub-url",
        required=True,
        help="URL of the AllSpice Hub instance to use.",
    )
    arg_parser.add_argument(
        "--log-level",
        required=False,
        default="INFO",
        help="Level of logger to use, default INFO.",
    )

    args = arg_parser.parse_args()
    logger.setLevel(args.log_level.upper())
    logger.info("Starting wireviz template generator, with args %s", vars(args))

    token_text = os.getenv("ALLSPICE_AUTH_TOKEN")
    if not token_text:
        raise ValueError("ALLSPICE_AUTH_TOKEN environment variable not set.")

    client = AllSpice(
        allspice_hub_url=args.allspice_hub_url,
        token_text=token_text,
        log_level=args.log_level.upper(),
    )
    owner_name, repo_name = args.repository.split("/")
    repository = Repository.request(client, owner_name, repo_name)

    assembly_data = repository.get_raw_file(args.assembly_file, ref=args.ref)
    assembly = MultiboardAssembly.load_yaml(io.StringIO(assembly_data.decode("utf-8")))
    variant = args.variant

    logger.info(
        "Loaded assembly file %s from repository %s",
        args.assembly_file,
        args.repository,
    )

    dependencies = list(assembly.dependencies)
    if variant:
        try:
            dependencies += assembly.variants[variant]
        except KeyError:
            raise ValueError(f"Variant {variant} not found in assembly file.")

    logger.info(
        "Found %d dependencies to generate wireviz template for",
        len(dependencies),
    )

    wireviz_connectors = {}
    for dependency in dependencies:
        components = get_dependency_components(client, dependency)
        connector_components = find_connector_components(components)

        for component in connector_components:
            connectors = convert_component_to_wireviz_connector(
                component,
                dependency,
            )
            for connector in connectors:
                wireviz_connectors |= connector.to_dict()

    wireviz_output = {"connectors": wireviz_connectors}

    with open(args.output_file, "w") as output_file:
        yaml.dump(wireviz_output, output_file)

    logger.info("Wireviz template generated and written to %s", args.output_file)


if __name__ == "__main__":
    main()