code/tt_diffusion_planner/tests/_research.py

# SPDX-License-Identifier: Apache-2.0
"""Access to the verified research scripts of the porting workspace (``research/diffusion-planner/scripts``), the
independent implementations the host code is unit-tested against. Absent in an installed package or the image:
tests that need them skip."""
from __future__ import annotations

import importlib.util
import sys
from pathlib import Path
from typing import Optional

import numpy as np

PKG = Path(__file__).resolve().parents[1]
RESEARCH = PKG.parents[3] / "research" / "diffusion-planner"
SCRIPTS = RESEARCH / "scripts"
ORT_GOLDENS = RESEARCH / "ort"
FULL_GOLDENS = RESEARCH / "goldens"
SAMPLES = PKG / "samples"
SMALL_GOLDENS = PKG / "tests" / "goldens"


def load_script(name: str):
    """Import ``research/diffusion-planner/scripts/<name>.py`` by path (None when absent). Its directory is put on
    ``sys.path`` only while it imports (``dp_scene`` imports ``dp_common`` by name)."""
    path = SCRIPTS / f"{name}.py"
    if not path.is_file():
        return None
    key = f"_dp_research_{name}"
    if key in sys.modules:
        return sys.modules[key]
    spec = importlib.util.spec_from_file_location(key, path)
    mod = importlib.util.module_from_spec(spec)
    sys.path.insert(0, str(SCRIPTS))
    try:
        spec.loader.exec_module(mod)
    finally:
        sys.path.remove(str(SCRIPTS))
    sys.modules[key] = mod
    return mod


def sample_raw(stem: str) -> dict:
    with np.load(SAMPLES / f"{stem}.npz") as z:
        return {k: np.array(z[k]) for k in z.files}


def research_scene_raw(scene: str) -> Optional[dict]:
    path = ORT_GOLDENS / f"golden_{scene}.npz"
    if not path.is_file():
        return None
    with np.load(path) as z:
        return {k[len("raw/"):]: np.array(z[k]) for k in z.files if k.startswith("raw/")}


def research_scenes() -> list:
    return sorted(p.stem[len("golden_"):] for p in ORT_GOLDENS.glob("golden_*.npz"))