mohan-bee/drone

work/extract_kicad_netlist.py

from __future__ import annotations

import json
import math
import sys
from collections import defaultdict
from pathlib import Path

import sexpdata

SOURCE = Path("/Users/mohan/Downloads/ESP32 drone.kicad_sch")


def tag(item):
    return str(item[0]) if isinstance(item, list) and item else ""


def children(item, name):
    return [x for x in item if tag(x) == name]


def child(item, name):
    return next((x for x in item if tag(x) == name), None)


def prop(item, name, default=""):
    p = next(
        (
            x
            for x in item
            if tag(x) == "property" and len(x) >= 3 and x[1] == name
        ),
        None,
    )
    return p[2] if p else default


def point(item):
    return (round(float(item[1]), 3), round(float(item[2]), 3))


class UnionFind:
    def __init__(self):
        self.parent = {}

    def add(self, x):
        self.parent.setdefault(x, x)

    def find(self, x):
        self.add(x)
        while self.parent[x] != x:
            self.parent[x] = self.parent[self.parent[x]]
            x = self.parent[x]
        return x

    def union(self, a, b):
        a, b = self.find(a), self.find(b)
        if a != b:
            self.parent[b] = a


data = sexpdata.load(SOURCE.open())
lib_symbols_node = child(data, "lib_symbols")
lib_symbols = {
    item[1]: item
    for item in lib_symbols_node[1:]
    if tag(item) == "symbol" and len(item) > 1
}


def pins_in_lib_symbol(node):
    found = {}

    def walk(item):
        if not isinstance(item, list):
            return
        if tag(item) == "pin":
            at = child(item, "at")
            number = child(item, "number")
            name = child(item, "name")
            if at and number:
                found[str(number[1])] = {
                    "name": str(name[1]) if name else "",
                    "at": (float(at[1]), float(at[2])),
                }
            return
        for sub in item:
            walk(sub)

    walk(node)
    return found


library_pins = {name: pins_in_lib_symbol(node) for name, node in lib_symbols.items()}
instances = [x for x in data if tag(x) == "symbol"]
wires = [x for x in data if tag(x) == "wire"]
labels = [x for x in data if tag(x) in ("label", "global_label")]
junctions = [point(child(x, "at")) for x in data if tag(x) == "junction"]

uf = UnionFind()
segments = []
for wire in wires:
    pts = child(wire, "pts")
    ends = children(pts, "xy")
    if len(ends) >= 2:
        a, b = point(ends[0]), point(ends[-1])
        uf.union(a, b)
        segments.append((a, b))


def on_segment(p, a, b, tolerance=0.002):
    cross = (p[0] - a[0]) * (b[1] - a[1]) - (p[1] - a[1]) * (b[0] - a[0])
    if abs(cross) > tolerance:
        return False
    return (
        min(a[0], b[0]) - tolerance <= p[0] <= max(a[0], b[0]) + tolerance
        and min(a[1], b[1]) - tolerance <= p[1] <= max(a[1], b[1]) + tolerance
    )


attachment_points = set(junctions)
for label in labels:
    attachment_points.add(point(child(label, "at")))

component_pins = []
power_pins = []
component_inventory = []

for instance in instances:
    lib_id = str(child(instance, "lib_id")[1])
    ref = str(prop(instance, "Reference"))
    value = str(prop(instance, "Value"))
    footprint = str(prop(instance, "Footprint"))
    at = child(instance, "at")
    ix, iy = float(at[1]), float(at[2])
    rotation = float(at[3]) if len(at) > 3 else 0
    mirror_node = child(instance, "mirror")
    mirror = str(mirror_node[1]) if mirror_node else ""
    rad = math.radians(rotation)
    instance_pin_numbers = {str(p[1]) for p in children(instance, "pin")}
    pin_defs = library_pins.get(lib_id, {})
    is_power = lib_id.startswith("power:")

    if not is_power:
        component_inventory.append(
            {
                "reference": ref,
                "value": value,
                "footprint": footprint,
                "lib_id": lib_id,
            }
        )

    for number in instance_pin_numbers:
        definition = pin_defs.get(number)
        if not definition:
            continue
        x, y = definition["at"]
        if mirror == "x":
            y = -y
        elif mirror == "y":
            x = -x
        gx = ix + math.cos(rad) * x + math.sin(rad) * y
        gy = iy + math.sin(rad) * x - math.cos(rad) * y
        coord = (round(gx, 3), round(gy, 3))
        attachment_points.add(coord)
        record = {
            "reference": ref,
            "value": value,
            "pin": number,
            "pin_name": definition["name"],
            "coord": coord,
        }
        if is_power:
            power_pins.append(record)
        else:
            component_pins.append(record)

# Join every explicit attachment point to each segment it lies on.
for p in attachment_points:
    uf.add(p)
    for a, b in segments:
        if on_segment(p, a, b):
            uf.union(p, a)

# Labels with the same name are electrically common on this single sheet.
label_coords = defaultdict(list)
for label in labels:
    label_coords[str(label[1])].append(point(child(label, "at")))
for coords in label_coords.values():
    for coord in coords[1:]:
        uf.union(coords[0], coord)

# KiCad power symbols of the same value are global.
power_coords = defaultdict(list)
for pin in power_pins:
    power_coords[pin["value"]].append(pin["coord"])
for coords in power_coords.values():
    for coord in coords[1:]:
        uf.union(coords[0], coord)

names_by_root = defaultdict(set)
for name, coords in label_coords.items():
    for coord in coords:
        names_by_root[uf.find(coord)].add(name)
for name, coords in power_coords.items():
    for coord in coords:
        names_by_root[uf.find(coord)].add(name)

pins_by_root = defaultdict(list)
for pin in component_pins:
    root = uf.find(pin["coord"])
    pins_by_root[root].append(pin)

nets = []
for root, pins in pins_by_root.items():
    names = sorted(names_by_root.get(root, []))
    auto_name = "~".join(
        f'{p["reference"]}.{p["pin"]}'
        for p in sorted(pins, key=lambda p: (p["reference"], p["pin"]))
    )
    nets.append(
        {
            "name": "/".join(names) if names else auto_name,
            "labels": names,
            "pins": sorted(pins, key=lambda p: (p["reference"], p["pin"])),
        }
    )

result = {
    "components": sorted(component_inventory, key=lambda c: c["reference"]),
    "nets": sorted(nets, key=lambda n: n["name"]),
}
json.dump(result, sys.stdout, indent=2)