instant-bom/bom_assistant/classification/classifier.py

from __future__ import annotations

import re

from bom_assistant.session.models import ComponentCategory

_PREFIX_MAP: dict[str, ComponentCategory] = {
    "C": ComponentCategory.capacitor,
    "R": ComponentCategory.resistor,
    "RV": ComponentCategory.resistor,
    "L": ComponentCategory.inductor,
    "D": ComponentCategory.diode,
    "U": ComponentCategory.ic,
    "IC": ComponentCategory.ic,
    "J": ComponentCategory.connector,
    "P": ComponentCategory.connector,
    "CON": ComponentCategory.connector,
    "USB": ComponentCategory.connector,
    "Q": ComponentCategory.mosfet,
    "T": ComponentCategory.mosfet,
    "X": ComponentCategory.crystal,
    "Y": ComponentCategory.crystal,
    "NTC": ComponentCategory.sensor,
    "RT": ComponentCategory.sensor,
}

_LED_KEYWORDS = {"red", "grn", "blu", "green", "blue", "led", "yellow", "white", "amber"}

_HW_RE = re.compile(
    r"^(PCB|PCBA|screw|washer|nut|lock\s*nut|heatsink|standoff|spacer|heat\s*sink)",
    re.IGNORECASE,
)

# Value-string heuristics (used when no designators)
_CAP_RE = re.compile(r"\d+\.?\d*\s*[pnuµ][fF]?$", re.IGNORECASE)
_RES_RE = re.compile(r"(\d+[kKMR]|\d+\.?\d*\s*[kKMΩR]$|\d+[kK]\d+)", re.IGNORECASE)
_PART_NUM_RE = re.compile(r"^[A-Z]{2,}[\dA-Z\-]{3,}$")


def classify(designators: list[str], raw_value: str) -> ComponentCategory:
    prefix = _extract_prefix(designators[0]) if designators else ""
    category = _PREFIX_MAP.get(prefix.upper())

    if category == ComponentCategory.diode:
        low = raw_value.lower()
        if any(kw in low for kw in _LED_KEYWORDS):
            return ComponentCategory.led
        return ComponentCategory.diode

    if category is not None:
        return category

    # no designator match — use value heuristics
    return _classify_by_value(raw_value)


def is_non_electronic(raw_value: str) -> bool:
    return bool(_HW_RE.match(raw_value.strip()))


def _extract_prefix(designator: str) -> str:
    m = re.match(r"^([A-Za-z]+)", designator.strip())
    return m.group(1) if m else ""


def _classify_by_value(raw_value: str) -> ComponentCategory:
    v = raw_value.strip()
    if _CAP_RE.match(v):
        return ComponentCategory.capacitor
    if _RES_RE.match(v):
        return ComponentCategory.resistor
    # looks like a part number: ≥5 chars, mixed alpha+digit, no unit suffix
    if _PART_NUM_RE.match(v) and len(v) >= 5:
        return ComponentCategory.ic
    return ComponentCategory.other