__init__.py

"""Auto-discover OpenMC model builders for benchmarks."""

from __future__ import annotations

import importlib
import pkgutil
from typing import Callable, Dict, Optional, Tuple

import openmc

ModelBuilder = Callable[[], openmc.Model]
CustomMetrics = Optional[Dict[str, Callable]]
ModelSpec = Tuple[str, ModelBuilder, Optional[Tuple[int, ...]], Optional[Tuple[Optional[int], ...]], CustomMetrics]

# Mapping of module name -> ModelSpec
MODEL_REGISTRY: Dict[str, ModelSpec] = {}


def _default_benchmark_name(module_name: str) -> str:
    parts = module_name.split("_")
    return "".join(part.capitalize() for part in parts)


def _discover() -> None:
    package = __name__
    for module_info in pkgutil.iter_modules(__path__):
        if module_info.name.startswith("_"):
            continue
        module = importlib.import_module(f"{package}.{module_info.name}")
        builder = getattr(module, "build_model", None)
        if builder is None:
            continue
        benchmark_name = getattr(module, "BENCHMARK_NAME", _default_benchmark_name(module_info.name))
        configs = getattr(module, "CONFIGS", None)
        custom_metrics = getattr(module, "CUSTOM_METRICS", None)
        MODEL_REGISTRY[module_info.name] = (benchmark_name, builder, configs, custom_metrics)


_discover()

for _module, (_benchmark_name, builder, _configs, _custom) in MODEL_REGISTRY.items():
    globals()[_module] = builder

__all__ = ['MODEL_REGISTRY', 'ModelBuilder', 'ModelSpec'] + sorted(MODEL_REGISTRY)