eng-lab/hwe/design.py

"""Load the design: every knob of the drone, from design/drone.toml, with units.

The design file is the single source of truth. Every analysis reads its inputs
from the Design object built here, so changing one number in drone.toml (or
overriding it with `python -m hwe tinker`) changes every downstream result on
the next run. That is the design loop of EL01, section 2.1, made mechanical.

Rules the loader enforces
- Every physical number is a string with a unit ("150 mm", "1300 mA*h").
  A bare number is accepted only for counts and pure ratios, which are named
  explicitly below. Anything else raises an error that names the key.
- Overrides use dotted keys: "battery.capacity=2200 mA*h".
"""
from __future__ import annotations

import copy
import re
import tomllib
from dataclasses import dataclass, field
from pathlib import Path

from hwe.units import Q_, Quantity, parse

ROOT = Path(__file__).resolve().parents[1]          # the eng-lab folder
DESIGN_FILE = ROOT / "design" / "drone.toml"

# A design-file string is a quantity exactly when it starts with a number.
_NUMBER_FIRST = re.compile(r"^\s*[-+]?(\d|\.\d)")

# Keys whose values are plain numbers on purpose (counts and ratios).
PLAIN_NUMBER_KEYS = {
    "propulsion.motor_count", "battery.cells", "battery.usable_fraction",
    "arm.count", "requirements.min_thrust_to_weight", "tests.phone.sign",
}


@dataclass
class Item:
    """One line of the mass budget: what hwe.core.massbudget receives."""
    name: str
    subsystem: str
    mass: Quantity
    maturity: str
    source: str = ""


@dataclass
class ComponentSpec:
    """One [[component]] entry as written in the design file.

    mass is the total for all `count` copies, or None when the mass comes from
    CAD (mass = "cad" in the file); then `fallback` (optional) is used until the
    CAD stage has run.
    """
    name: str
    subsystem: str
    mass: Quantity | None
    maturity: str
    source: str = ""
    count: int = 1
    fallback: Quantity | None = None


@dataclass
class Design:
    """The whole design as nested dictionaries of Quantities, plus the parts list."""
    raw: dict                       # the TOML as read, after overrides (strings)
    values: dict                    # same shape, with units parsed into Quantities
    items: list[ComponentSpec] = field(default_factory=list)

    def get(self, dotted: str):
        """design.get("battery.capacity") -> Q_(1300, "mA*h")"""
        node = self.values
        for part in dotted.split("."):
            node = node[part]
        return node

    def section(self, name: str) -> dict:
        return self.values.get(name, {})


def _parse_value(dotted: str, value):
    if isinstance(value, dict):
        return {k: _parse_value(f"{dotted}.{k}", v) for k, v in value.items()}
    if isinstance(value, list):
        return value
    if isinstance(value, bool):
        return value
    if isinstance(value, (int, float)):
        if dotted in PLAIN_NUMBER_KEYS or dotted.startswith("growth_allowance."):
            return value
        raise TypeError(f"{dotted} = {value!r} has no unit; write it as a string such as \"{value} g\"")
    if isinstance(value, str):
        if _NUMBER_FIRST.match(value):  # "150 mm", "-3 degC", ".5 kg": a quantity
            return parse(value)         # a typo in the unit raises here, naming the unit
        return value                    # names, paths, notes, maturities stay text
    return value


def _set_dotted(tree: dict, dotted: str, text: str) -> None:
    keys = dotted.split(".")
    node = tree
    for k in keys[:-1]:
        if k not in node:
            raise KeyError(f"override {dotted}: no section '{k}' in the design file")
        node = node[k]
    if keys[-1] not in node:
        raise KeyError(f"override {dotted}: no key '{keys[-1]}' in [{'.'.join(keys[:-1])}]")
    old = node[keys[-1]]
    if isinstance(old, (int, float)) and not isinstance(old, bool):
        node[keys[-1]] = type(old)(text)
    else:
        node[keys[-1]] = text


def load(path: Path | str = DESIGN_FILE, overrides: dict[str, str] | None = None) -> Design:
    """Read the design file, apply overrides such as {"battery.capacity": "2200 mA*h"}."""
    with open(path, "rb") as f:
        raw = tomllib.load(f)
    raw = copy.deepcopy(raw)
    for dotted, text in (overrides or {}).items():
        _set_dotted(raw, dotted, text)

    components = raw.get("component", [])
    body = {k: v for k, v in raw.items() if k != "component"}
    values = {k: _parse_value(k, v) for k, v in body.items()}

    items = []
    for entry in components:
        count = entry.get("count", 1)
        if isinstance(count, str):              # count = "propulsion.motor_count" follows that knob
            node = values
            for part in count.split("."):
                node = node[part]
            count = node
        count = int(count)
        from_cad = entry.get("mass") == "cad"
        items.append(ComponentSpec(
            name=entry["name"] if count == 1 else f"{entry['name']} (x{count})",
            subsystem=entry["subsystem"],
            mass=None if from_cad else parse(entry["mass"]) * count,
            maturity=entry["maturity"],
            source=entry.get("source", ""),
            count=count,
            fallback=parse(entry["fallback"]) * count if "fallback" in entry else None,
        ))
    return Design(raw=raw, values=values, items=items)


def parse_override(text: str) -> tuple[str, str]:
    """'battery.capacity=2200 mA*h' -> ('battery.capacity', '2200 mA*h')"""
    if "=" not in text:
        raise ValueError(f"override must look like key=value, got {text!r}")
    key, value = text.split("=", 1)
    return key.strip(), value.strip()