imrishabh18/pcbgolf

footprints/generate.py

#!/usr/bin/env python3
"""Convert the challenge's own KiCad footprints, retaining exact copper geometry.

Usage: python footprints/generate.py /path/to/PCBGolf/pcbgolf.pretty
Requires the installed `tsci` CLI. Its graphics conversion is retained; pads are
regenerated from the KiCad source to retain radii, mask/paste margins, and angles.
"""
import concurrent.futures
import json
import math
from pathlib import Path
import re
import subprocess
import sys
import tempfile

DEST = Path(__file__).resolve().parent
SOURCE = Path(sys.argv[1]).resolve()
# Corrections retain their own native source footprint alongside this generator.
ADDITIONAL_SOURCE = DEST / 'kicad'
INSERTION_DIRECTIONS = {
    'DCJACK_2MM_SMT': 'from_left',
    'DX07S024XJ1R1100': 'from_bottom',
    'USB_C_Receptacle_GCT_USB4105-xx-A_16P_TopMnt_Horizontal': 'from_bottom',
    '0472192001': 'from_above',
    '2X04': 'from_above',
    'USB-C-FEMALE-VERT-GCT': 'from_above',
}


def parse_sexpr(text):
    tokens = re.findall(r'"(?:\\.|[^"\\])*"|[()]|[^\s()]+', text)
    stack, root = [], None
    for token in tokens:
        if token == '(':
            node = []
            if stack:
                stack[-1].append(node)
            else:
                root = node
            stack.append(node)
        elif token == ')':
            stack.pop()
        else:
            stack[-1].append(json.loads(token) if token.startswith('"') else token)
    return root


def child(node, name, default=None):
    return next((n for n in node if isinstance(n, list) and n[0] == name), default)


def value(node, name, default=None):
    item = child(node, name)
    return item[1] if item else default


def number(value):
    return round(float(value), 9)


def attrs(props):
    return ' '.join(f'{key}={{{json.dumps(val, separators=(",", ":"))}}}' for key, val in props.items())


def convert(path):
    source = parse_sexpr(path.read_text())
    name = path.stem
    pads = [n for n in source if isinstance(n, list) and n[0] == 'pad']
    pin_ids = list(dict.fromkeys(p[1] for p in pads if p[1] and p[2] != 'np_thru_hole'))
    pin_map = {p: f'pin{int(p)}' for p in pin_ids if p.isdigit() and int(p) > 0}
    last_pin = max([int(p[3:]) for p in pin_map.values()] or [0])
    for pad_id in pin_ids:
        if pad_id not in pin_map:
            last_pin += 1
            pin_map[pad_id] = f'pin{last_pin}'

    with tempfile.TemporaryDirectory() as temp:
        output = Path(temp) / 'converted.tsx'
        subprocess.run(['tsci', 'convert', str(path), '--name', 'SourceFootprint', '-o', str(output)], check=True, capture_output=True, text=True)
        converted = output.read_text()
    graphics = converted.split('<footprint>', 1)[1].split('</footprint>', 1)[0]
    graphics = re.sub(r'<(?:smtpad|platedhole|hole)\b[^>]*?/>', '', graphics, flags=re.S)
    # Component reference/value labels are provided by the board, not REF**.
    graphics = re.sub(r'<(?:silkscreentext|fabricationnotetext)\b[^>]*(?:REF\*\*|reference)[^>]*?/>', '', graphics, flags=re.S | re.I)
    graphics_lines = [line.strip() for line in graphics.splitlines() if line.strip()]
    # Assembly courtyards are reviewed against native CAD bodies and pad extents.
    courtyard_path = DEST / 'courtyard-additions.json'
    courtyards = json.loads(courtyard_path.read_text()) if courtyard_path.exists() else {}
    if name in courtyards:
        row = courtyards[name]
        outline = row.get('outline', [[row['xmin'], row['ymin']], [row['xmax'], row['ymin']], [row['xmax'], row['ymax']], [row['xmin'], row['ymax']]])
        graphics_lines = [line for line in graphics_lines if '<courtyard' not in line]
        points = [{'x': round(x, 6), 'y': round(y, 6)} for x, y in outline]
        if row.get('shape') == 'rect':
            cx = round((row['xmin'] + row['xmax']) / 2, 6)
            cy = round((row['ymin'] + row['ymax']) / 2, 6)
            width = round(row['xmax'] - row['xmin'], 6)
            height = round(row['ymax'] - row['ymin'], 6)
            graphics_lines.append(f'<courtyardrect pcbX={{{cx}}} pcbY={{{cy}}} width={{{width}}} height={{{height}}} layer="top" />')
        else:
            graphics_lines.append('<courtyardoutline outline={' + json.dumps(points) + '} layer="top" />')
    # A collection of two-point outlines does not describe a closed area in
    # Circuit JSON. Join connected source courtyard edges into closed polygons.
    edges = []
    retained = []
    for line in graphics_lines:
        match = re.search(r'<courtyardoutline outline=\{(\[.*?\])\}', line)
        points = json.loads(match[1]) if match else []
        if len(points) == 2:
            edges.append(points)
        else:
            retained.append(line)
    while edges:
        outline = edges.pop(0)
        while outline[-1] != outline[0]:
            for i, edge in enumerate(edges):
                if edge[0] == outline[-1] or edge[1] == outline[-1]:
                    outline.append(edge[1] if edge[0] == outline[-1] else edge[0])
                    edges.pop(i)
                    break
            else:
                raise ValueError('Open source courtyard: ' + name)
        retained.append('<courtyardoutline outline={' + json.dumps(outline[:-1]) + '} layer="top" />')
    graphics_lines = retained

    copper_lines, copper_metadata = [], []
    for index, pad in enumerate(pads):
        pad_id, kind, shape = pad[1:4]
        at = child(pad, 'at')
        size = child(pad, 'size')
        x, y = number(at[1]), -number(at[2])
        rotation = number(at[3]) % 360 if len(at) > 3 else 0
        width, height = number(size[1]), number(size[2])
        props = {'name': f'pad_{index + 1}', 'pcbX': x, 'pcbY': y}
        if kind != 'np_thru_hole' and pad_id:
            # The canonical pin hint creates a robust generic-chip interface;
            # the original pad ID remains available for inspection/selectors.
            props['portHints'] = [pin_map[pad_id], pad_id]
        layers = child(pad, 'layers')[1:]
        mask = value(pad, 'solder_mask_margin', value(source, 'solder_mask_margin', 0))
        paste = value(pad, 'solder_paste_margin', value(source, 'solder_paste_margin', 0))
        if mask is not None:
            props['solderMaskMargin'] = number(mask)
        if kind == 'smd' and paste is not None:
            props['solderPasteMargin'] = number(paste)
        if kind == 'smd':
            props['layer'] = 'bottom' if 'B.Cu' in layers else 'top'
            assert 'F.Cu' in layers or 'B.Cu' in layers, (name, pad_id, layers)
            radius = min(width, height) * number(value(pad, 'roundrect_rratio', 0))
            # A square roundrect with a half-width corner radius is exactly a
            # circle. Use the equivalent circle primitive so clearance checks
            # measure copper rather than the square enclosing that copper.
            if shape == 'roundrect' and abs(width - height) < 1e-9 and abs(radius - width / 2) < 1e-9:
                props.update(shape='circle', radius=width / 2)
            elif rotation % 90 == 0:
                if rotation % 180:
                    width, height = height, width
                if shape in ['rect', 'roundrect']:
                    props.update(shape='rect', width=width, height=height)
                    if radius:
                        props['cornerRadius'] = number(radius)
                elif shape == 'circle':
                    assert width == height
                    props.update(shape='circle', radius=width / 2)
                elif shape == 'oval':
                    props.update(shape='pill', width=width, height=height, radius=min(width, height) / 2)
                else:
                    raise ValueError((name, pad_id, shape))
            else:
                assert shape in ['rect', 'roundrect'], (name, pad_id, shape)
                props.update(shape='rotated_rect', width=width, height=height, ccwRotation=rotation)
                if radius:
                    props['cornerRadius'] = number(radius)
            tag = 'smtpad'
        elif kind == 'np_thru_hole':
            drill = child(pad, 'drill')
            if drill[1] == 'oval':
                props.update(shape='pill', width=number(drill[2]), height=number(drill[3]), pcbRotation=rotation)
            else:
                props['diameter'] = number(drill[1])
            tag = 'hole'
        elif kind == 'thru_hole':
            drill = child(pad, 'drill')
            if drill[1] == 'oval':
                if shape in ['rect', 'roundrect']:
                    props.update(shape='pill_hole_with_rect_pad', holeWidth=number(drill[2]), holeHeight=number(drill[3]), rectPadWidth=width, rectPadHeight=height, pcbRotation=rotation)
                else:
                    props.update(shape='pill', holeWidth=number(drill[2]), holeHeight=number(drill[3]), outerWidth=width, outerHeight=height, pcbRotation=rotation)
            elif shape == 'circle':
                assert width == height, (name, pad_id, size)
                props.update(shape='circle', outerDiameter=width, holeDiameter=number(drill[1]))
            elif shape in ['rect', 'roundrect']:
                props.update(shape='circular_hole_with_rect_pad', holeDiameter=number(drill[1]), rectPadWidth=width, rectPadHeight=height, pcbRotation=rotation)
                if shape == 'roundrect':
                    props['rectBorderRadius'] = number(min(width, height) * number(value(pad, 'roundrect_rratio', 0)))
            else:
                raise ValueError((name, pad_id, shape))
            tag = 'platedhole'
        else:
            raise ValueError((name, pad_id, kind))
        copper_lines.append(f'<{tag} {attrs(props)} />')
        copper_metadata.append({'sourcePadId': pad_id, 'sourceKind': kind, 'sourceShape': shape, 'element': tag, **props})

    fn = 'Footprint_' + re.sub(r'[^A-Za-z0-9_]', '_', name)
    filename = fn + '.tsx'
    body = '\n'.join('    ' + line for line in copper_lines + graphics_lines)
    insertion = f' insertionDirection="{INSERTION_DIRECTIONS[name]}"' if name in INSERTION_DIRECTIONS else ''
    origin = f'footprints/kicad/{path.name} (see source-corrections.json)' if path.parent == ADDITIONAL_SOURCE else f'commaai/PCBGolf/pcbgolf.pretty/{path.name}'
    (DEST / filename).write_text(f'// Source: {origin}\n// Generated by footprints/generate.py. Coordinates are millimetres, +Y up.\nexport const {fn} = () => (\n  <footprint originalLayer="top"{insertion}>\n{body}\n  </footprint>\n)\n')
    return name, filename, fn, pin_map, copper_metadata


files = sorted(SOURCE.glob('*.kicad_mod')) + sorted(ADDITIONAL_SOURCE.glob('*.kicad_mod'))
assert len({p.stem for p in files}) == len(files), 'Duplicate footprint names across original and correction libraries'
assert files, f'No .kicad_mod files found in {SOURCE}'
with concurrent.futures.ThreadPoolExecutor(max_workers=4) as pool:
    results = list(pool.map(convert, files))
imports = '\n'.join(f'import {{ {fn} }} from "./{filename[:-4]}"' for name, filename, fn, pins, pads in results)
mapping = ',\n'.join(f'  {json.dumps(name)}: {fn}' for name, filename, fn, pins, pads in results)
pins = {name: pin_map for name, filename, fn, pin_map, pads in results}
(DEST / 'index.tsx').write_text(imports + '\n\n/** Exact challenge footprints, keyed by original KiCad library name. */\nexport const footprints = {\n' + mapping + '\n} as const\n\n/** Original KiCad pad ID -> tscircuit canonical pin hint. */\nexport const footprintPinMaps: Record<string, Record<string, string>> = ' + json.dumps(pins, indent=2) + '\n\nexport type FootprintName = keyof typeof footprints\n')
metadata = {name: {'pinMap': pin_map, 'pads': pads} for name, filename, fn, pin_map, pads in results}
(DEST / 'geometry.json').write_text(json.dumps(metadata, indent=2) + '\n')
print(f'Generated {len(results)} footprints, {sum(len(pads) for *_, pads in results)} pads/holes.')