agent_threat_map/loader.py

from __future__ import annotations

import json
from pathlib import Path

from agent_threat_map.schema import Probe


def load_probes(path: str | Path) -> list[Probe]:
    path = Path(path)
    probes: list[Probe] = []
    with path.open(encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            probes.append(Probe.from_dict(json.loads(line)))
    return probes


def load_categories(path: str | Path) -> dict:
    with Path(path).open(encoding="utf-8") as f:
        return json.load(f)