1"""Thin wrapper around the `xctrace` CLI."""
2from __future__ import annotations
3 
4import subprocess
5import xml.etree.ElementTree as ET
6from dataclasses import dataclass
7from pathlib import Path
8 
9 
10@dataclass(frozen=True)
11class RunInfo:
12    """Per-run metadata and schemas. Instruments traces can hold multiple runs."""
13    number: int
14    template_name: str | None
15    duration_s: float | None
16    start_date: str | None
17    end_date: str | None
18    schemas: frozenset[str]
19 
20 
21@dataclass(frozen=True)
22class TraceInfo:
23    xctrace_version: str
24    runs: tuple[RunInfo, ...]
25 
26    def get_run(self, number: int) -> RunInfo:
27        for r in self.runs:
28            if r.number == number:
29                return r
30        available = ", ".join(str(r.number) for r in self.runs)
31        raise KeyError(f"run {number} not in trace (available: {available})")
32 
33 
34def version() -> str:
35    out = subprocess.run(
36        ["xctrace", "version"], capture_output=True, text=True, check=True
37    )
38    return out.stdout.strip()
39 
40 
41def toc(trace_path: Path) -> TraceInfo:
42    """Export the trace's table of contents and return per-run metadata.
43 
44    The TOC is small (a few KB) so we load it fully rather than streaming.
45    """
46    xml_bytes = _run_export(trace_path, ["--toc"])
47    root = ET.fromstring(xml_bytes)
48 
49    instruments = _find_text(root, ".//instruments-version") or ""
50 
51    runs: list[RunInfo] = []
52    for run_el in root.iterfind("./run"):
53        number_attr = run_el.get("number")
54        if not number_attr:
55            continue
56        try:
57            number = int(number_attr)
58        except ValueError:
59            continue
60        if number <= 0:
61            continue
62 
63        schemas: set[str] = set()
64        for table in run_el.iterfind("./data/table"):
65            schema = table.get("schema")
66            if schema:
67                schemas.add(schema)
68 
69        summary = run_el.find("./info/summary")
70        if summary is not None:
71            template = _find_text(summary, "./template-name")
72            duration = _find_text(summary, "./duration")
73            start = _find_text(summary, "./start-date")
74            end = _find_text(summary, "./end-date")
75        else:
76            template = duration = start = end = None
77 
78        runs.append(RunInfo(
79            number=number,
80            template_name=template,
81            duration_s=float(duration) if duration else None,
82            start_date=start,
83            end_date=end,
84            schemas=frozenset(schemas),
85        ))
86 
87    runs.sort(key=lambda r: r.number)
88    return TraceInfo(
89        xctrace_version=instruments,
90        runs=tuple(runs),
91    )
92 
93 
94def export_schema(trace_path: Path, schema: str, run: int = 1) -> bytes:
95    """Export one schema's data as XML bytes from the given run.
96 
97    Callers are expected to iterparse the result rather than build a full tree
98    for large schemas (time-profile can be tens of MB).
99    """
100    xpath = f'/trace-toc/run[@number="{run}"]/data/table[@schema="{schema}"]'
101    return _run_export(trace_path, ["--xpath", xpath])
102 
103 
104def _run_export(trace_path: Path, extra_args: list[str]) -> bytes:
105    cmd = ["xctrace", "export", "--input", str(trace_path), *extra_args]
106    proc = subprocess.run(cmd, capture_output=True, check=False)
107    if proc.returncode != 0:
108        raise RuntimeError(
109            f"xctrace export failed ({proc.returncode}): "
110            f"{proc.stderr.decode(errors='replace').strip()}"
111        )
112    return proc.stdout
113 
114 
115def _find_text(root: ET.Element, path: str) -> str | None:
116    el = root.find(path)
117    return el.text if el is not None and el.text else None
118