traceforge-pcb 1.0.0__py3-none-any.whl

traceforge/tools/batch.py

@@ -0,0 +1,155 @@
+"""Run the tool calls of one model step concurrently where that is safe.
+
+A model step can ask for several tools at once. Each call gets a scheduling class:
+
+* ``read``: changes nothing (file reads, searches, KiCad queries, read-only shell
+  commands). Reads run concurrently, at most ``MAX_CONCURRENT_READS`` at a time.
+* ``write``: changes one known file. It waits for every earlier call that touches the
+  same file (or an unknown set of files), so edits apply in the order the model gave.
+* ``exclusive``: anything else (other shell commands, questions to the user, unknown
+  tools, writes whose target is unknown). It waits for every earlier call, and every
+  later call waits for it.
+
+Results are returned in the order of the calls, whatever order they finish in.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import os
+import shlex
+from collections.abc import Awaitable, Callable, Sequence
+from dataclasses import dataclass
+from typing import Any, TypeVar
+
+MAX_CONCURRENT_READS = 8
+
+READ = "read"
+WRITE = "write"
+EXCLUSIVE = "exclusive"
+
+
+@dataclass(frozen=True)
+class CallPlan:
+    kind: str  # read | write | exclusive
+    # Normalised file the call reads or writes; None means "not scoped to one file".
+    path: str | None = None
+
+
+def path_key(raw: Any, workspace: str | None = None) -> str | None:
+    if not isinstance(raw, str) or not raw.strip():
+        return None
+    path = raw.strip().replace("\\", "/")
+    if not os.path.isabs(path) and workspace:
+        path = os.path.join(workspace, path)
+    return os.path.normpath(path)
+
+
+_READ_PROGRAMS = {
+    "ls",
+    "cat",
+    "head",
+    "wc",
+    "grep",
+    "rg",
+    "find",
+    "pwd",
+    "stat",
+    "file",
+    "du",
+    "diff",
+    "which",
+}
+_FIND_ACTIONS = {"-exec", "-execdir", "-ok", "-okdir", "-delete", "-fprint", "-fprintf", "-fls"}
+
+
+def read_only_command(command: Any) -> bool:
+    """Conservatively decide that a shell command only reads."""
+    if not isinstance(command, str) or not command.strip():
+        return False
+    if any(token in command for token in ("\n", ";", "&", "|", "`", "<", ">", "$(", "${")):
+        return False
+    try:
+        words = shlex.split(command)
+    except ValueError:
+        return False
+    if not words or os.path.basename(words[0]) not in _READ_PROGRAMS:
+        return False
+    program, rest = os.path.basename(words[0]), words[1:]
+    if program == "find":
+        return not any(word in _FIND_ACTIONS for word in rest)
+    if program == "rg":
+        return not any(word.startswith("--pre") for word in rest)
+    return True
+
+
+def _overlaps(left: CallPlan, right: CallPlan) -> bool:
+    if WRITE not in (left.kind, right.kind):
+        return False
+    if left.path is None or right.path is None:
+        return True
+    return left.path == right.path
+
+
+def dependencies(plans: Sequence[CallPlan]) -> list[set[int]]:
+    """For each call, the indexes of earlier calls it has to wait for."""
+    waits: list[set[int]] = []
+    last_exclusive: int | None = None
+    for index, plan in enumerate(plans):
+        mine: set[int] = set()
+        if plan.kind == EXCLUSIVE:
+            mine.update(range(index))
+            last_exclusive = index
+        else:
+            if last_exclusive is not None:
+                mine.add(last_exclusive)
+            mine.update(i for i in range(index) if _overlaps(plans[i], plan))
+        waits.append(mine)
+    return waits
+
+
+T = TypeVar("T")
+R = TypeVar("R")
+
+
+async def run_batch(
+    calls: Sequence[T],
+    plans: Sequence[CallPlan],
+    handler: Callable[[T], Awaitable[R]],
+    *,
+    max_concurrent_reads: int = MAX_CONCURRENT_READS,
+) -> list[R]:
+    """Run ``handler`` on every call under the rules above; results in call order.
+
+    ``handler`` must not raise for an ordinary tool failure (return an error result
+    instead). If it does raise, the remaining calls are cancelled and the error
+    propagates. Cancelling the batch cancels every call still running.
+    """
+    if len(calls) != len(plans):
+        raise ValueError("one plan per call")
+    waits = dependencies(plans)
+    finished = [asyncio.Event() for _ in calls]
+    reads = asyncio.Semaphore(max(1, max_concurrent_reads))
+    results: list[Any] = [None] * len(calls)
+
+    async def one(index: int) -> None:
+        try:
+            for earlier in sorted(waits[index]):
+                await finished[earlier].wait()
+            if plans[index].kind == READ:
+                async with reads:
+                    results[index] = await handler(calls[index])
+            else:
+                results[index] = await handler(calls[index])
+        finally:
+            finished[index].set()
+
+    tasks = [asyncio.create_task(one(index)) for index in range(len(calls))]
+    try:
+        await asyncio.gather(*tasks)
+    finally:
+        for task in tasks:
+            if not task.done():
+                task.cancel()
+        await asyncio.gather(*tasks, return_exceptions=True)
+    return results

traceforge/tools/connections.py

@@ -0,0 +1,245 @@
+"""``kicad_connections``: read-only connectivity questions about the design.
+
+* ``pin`` (``U1.3`` or ``U1.VCC``): the pin's net and every other pin on it, with
+  pin names and electrical types.
+* ``part`` (``U1``): each pin of the part and its net.
+* ``net``: every pin on a net.
+* ``from`` + ``to`` (pins or parts): shortest paths between them, walking from net
+  to net through two-terminal passives (R, C, L, ferrite, fuse, diode, jumper) but
+  not through ICs, connectors, ground or supply rails (which would connect everything).
+
+Connectivity comes from the schematic netlist (kicad-cli) when available, else from
+the board's pad nets; pin names from the schematic symbols, else the pads' functions.
+"""
+
+from __future__ import annotations
+
+from collections import deque
+from typing import Any
+
+from traceforge.tools.base import Tool, ToolContext, ToolError
+
+_PASSIVE = {"R", "C", "L", "FB", "F", "D", "JP", "SJ", "RN"}
+_MAX_NODES = 60
+
+
+def _model(args: dict[str, Any], ctx: ToolContext):
+    from traceforge.kicad.cli import KicadCli, KicadCliError
+    from traceforge.kicad.pcb import read_board
+    from traceforge.kicad.project import discover
+    from traceforge.kicad.schematic import read_schematic
+    from traceforge.validate.netmodel import build_model
+
+    ws = discover(ctx.workspace)
+    sch_path = ctx.resolve(args["schematic"], must_exist=True) if args.get("schematic") else None
+    pcb_path = ctx.resolve(args["pcb"], must_exist=True) if args.get("pcb") else None
+    sch_path, pcb_path = sch_path or ws.schematic, pcb_path or ws.pcb
+    if sch_path is None and pcb_path is None:
+        if ws.others.get("altium"):  # Altium project: use the converted model
+            from traceforge.altium import AltiumError, load_project
+            from traceforge.altium.model import to_design_model
+
+            try:
+                return to_design_model(load_project(ctx.workspace))
+            except AltiumError as exc:
+                raise ToolError(str(exc)) from exc
+        raise ToolError("no .kicad_sch or .kicad_pcb in the workspace")
+    source = args.get("source") or "auto"
+    schematic = read_schematic(sch_path) if sch_path else None
+    board = read_board(pcb_path) if pcb_path else None
+    netlist = None
+    if schematic is not None and source in {"auto", "schematic"}:
+        cli = KicadCli()
+        if cli.available:
+            try:
+                netlist = cli.netlist(sch_path, ctx.workspace / ".traceforge" / "reports")
+            except KicadCliError:
+                netlist = None
+        if netlist is None and (source == "schematic" or board is None):
+            raise ToolError("schematic connectivity needs kicad-cli; pass source=board")
+    if source == "board" and board is None:
+        raise ToolError("no board found")
+    return build_model(schematic, board, netlist)
+
+
+def _pin_info(model, ref: str, pin: str) -> dict[str, Any]:
+    part = model.parts.get(ref)
+    info: dict[str, Any] = {"pin": f"{ref}.{pin}"}
+    definition = part.pin_defs.get(pin) if part else None
+    name = definition.name if definition else ""
+    if not name:
+        name = next((p.function for p in model.pads(ref, pin) if p.function), "")
+    if name and name not in {pin, "~"}:
+        info["name"] = name.replace("~{", "").replace("}", "")
+    if definition and definition.type:
+        info["type"] = definition.type
+    if part is not None and part.value:
+        info["value"] = part.value
+    return info
+
+
+def _resolve_pin(model, text: str) -> tuple[str, str]:
+    ref, sep, pin = text.strip().replace(":", ".").partition(".")
+    if not sep or not pin:
+        raise ToolError(f"{text!r}: give a pin as REF.PIN, e.g. U1.3 or U1.VCC")
+    if ref not in model.parts and not model.part_nets(ref):
+        raise ToolError(f"no part {ref}")
+    nets = model.part_nets(ref)
+    part = model.parts.get(ref)
+    if pin in nets or (part and pin in part.pin_defs):
+        return ref, pin
+    wanted = pin.upper()
+    matches = []
+    for number in {*nets, *(part.pin_defs if part else {})}:
+        name = _pin_info(model, ref, number).get("name", "")
+        if name.upper() == wanted or wanted in name.upper().split("/"):
+            matches.append(number)
+    if len(matches) == 1:
+        return ref, matches[0]
+    if matches:
+        raise ToolError(f"{text}: several pins match ({', '.join(sorted(matches))}); use a number")
+    raise ToolError(f"{ref} has no pin {pin}; pins: {', '.join(sorted(nets)) or 'none connected'}")
+
+
+def _is_rail(net: str) -> bool:
+    from traceforge.validate.circuits import is_ground, rail_voltage
+
+    return is_ground(net) or rail_voltage(net) is not None
+
+
+def _net_nodes(model, net: str, exclude: tuple[str, str] | None = None) -> list[dict]:
+    nodes = [_pin_info(model, r, p) for r, p in model.nodes(net) if (r, p) != exclude]
+    return nodes[:_MAX_NODES]
+
+
+def _find_net(model, name: str) -> str:
+    wanted = name.strip().lstrip("/").upper()
+    for net in model.nets:
+        if net.lstrip("/").upper() == wanted:
+            return net
+    close = [n for n in model.nets if wanted in n.upper()][:10]
+    raise ToolError(f"no net {name}" + (f"; similar: {', '.join(close)}" if close else ""))
+
+
+def _endpoints(model, text: str) -> tuple[set[str], set[str]]:
+    """(start nets, owning refs) for a pin, a whole part or a net."""
+    if "." in text or ":" in text:
+        ref, pin = _resolve_pin(model, text)
+        net = model.net_of(ref, pin)
+        return ({net} if net else set()), {ref}
+    nets = set(model.part_nets(text).values())
+    if not nets:
+        return {_find_net(model, text)}, set()
+    signal = {n for n in nets if not _is_rail(n)}
+    return (signal or nets), {text}  # a shared GND/supply says nothing about a path
+
+
+def _paths(model, src: str, dst: str, max_hops: int) -> list[list[str]]:
+    starts, src_refs = _endpoints(model, src)
+    targets, dst_refs = _endpoints(model, dst)
+    if starts & targets:
+        net = sorted(starts & targets)[0]
+        return [[src, f"net {net}", dst]]
+    # BFS over nets; a step crosses one two-terminal passive to its other net.
+    queue = deque((net, [f"net {net}"]) for net in sorted(starts))
+    seen = set(starts)
+    found: list[list[str]] = []
+    while queue and len(found) < 3:
+        net, path = queue.popleft()
+        if len(path) // 2 >= max_hops:
+            continue
+        for ref, pin in model.nodes(net):
+            part = model.parts.get(ref)
+            if ref in src_refs or ref in dst_refs or part is None or part.prefix not in _PASSIVE:
+                continue
+            pins = model.part_nets(ref)
+            if len(pins) != 2:
+                continue
+            for other_pin, other in pins.items():
+                if other_pin == pin or other in seen:
+                    continue
+                step = [*path, f"{ref} ({part.value}) pin {pin}->{other_pin}", f"net {other}"]
+                if other in targets:
+                    found.append([src, *step, dst])
+                    continue
+                if _is_rail(other):
+                    continue  # do not wander through GND/supply rails
+                seen.add(other)
+                queue.append((other, step))
+    return found
+
+
+def kicad_connections(args: dict[str, Any], ctx: ToolContext) -> dict:
+    model = _model(args, ctx)
+    out: dict[str, Any] = {"source": model.source}
+    asked = False
+    if args.get("pin"):
+        asked = True
+        ref, pin = _resolve_pin(model, str(args["pin"]))
+        net = model.net_of(ref, pin)
+        out["pin"] = _pin_info(model, ref, pin)
+        out["net"] = net or "(unconnected)"
+        if net:
+            out["connected_to"] = _net_nodes(model, net, (ref, pin))
+    if args.get("part"):
+        asked = True
+        ref = str(args["part"]).strip()
+        nets = model.part_nets(ref)
+        if ref not in model.parts and not nets:
+            raise ToolError(f"no part {ref}")
+        part = model.parts.get(ref)
+        out["part"] = {"ref": ref, "value": part.value if part else "", "pins": []}
+        numbers = sorted({*nets, *(part.pin_defs if part else {})}, key=_pin_key)
+        for number in numbers:
+            row = _pin_info(model, ref, number)
+            net = nets.get(number)
+            row["net"] = net or "(unconnected)"
+            if net:
+                row["others"] = len(model.nodes(net)) - 1
+            out["part"]["pins"].append(row)
+    if args.get("net"):
+        asked = True
+        net = _find_net(model, str(args["net"]))
+        out["net_pins"] = {"net": net, "pins": _net_nodes(model, net)}
+    if args.get("from") or args.get("to"):
+        asked = True
+        if not (args.get("from") and args.get("to")):
+            raise ToolError("a path needs both 'from' and 'to'")
+        paths = _paths(model, str(args["from"]), str(args["to"]), int(args.get("max_hops") or 4))
+        out["paths"] = [" -> ".join(p) for p in paths]
+        if not paths:
+            out["note"] = (
+                "no path through passives within max_hops (ICs, connectors, GND and supply "
+                "rails are not traversed)"
+            )
+    if not asked:
+        raise ToolError("ask about a 'pin' (U1.3), a 'part' (U1), a 'net', or 'from' + 'to'")
+    return out
+
+
+def _pin_key(number: str) -> tuple[int, str]:
+    return (int(number), "") if number.isdigit() else (10**6, number)
+
+
+CONNECTIONS_TOOL = Tool(
+    "kicad_connections",
+    "Connectivity questions, read-only: pin='U1.3' or 'U1.VCC' -> its net and every other "
+    "pin on it (names, types); part='U1' -> each pin and its net; net='+3V3' -> its pins; "
+    "from='U2.PB3' to='D1' -> paths through passives (R/C/L/D/fuse; not ICs or rails). "
+    "Uses the schematic netlist (kicad-cli) when available, else the board.",
+    {
+        "type": "object",
+        "properties": {
+            "pin": {"type": "string"},
+            "part": {"type": "string"},
+            "net": {"type": "string"},
+            "from": {"type": "string"},
+            "to": {"type": "string"},
+            "max_hops": {"type": "integer", "description": "passives to cross (default 4)"},
+            "source": {"type": "string", "enum": ["auto", "schematic", "board"]},
+            "schematic": {"type": "string"},
+            "pcb": {"type": "string"},
+        },
+    },
+    kicad_connections,
+)