File size: 1,508 Bytes
d766458 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | import os
from importlib import import_module
from typing import Callable, Optional, cast
from pxdbench.tools.ptx.interface import ProtenixAPI
PtxFactory = Callable[..., ProtenixAPI]
_FACTORIES: dict[str, PtxFactory] = {}
def register(name: str, factory: PtxFactory) -> None:
_FACTORIES[name] = factory
def _load_from_dotted(mod_path: str, cls_name: str) -> PtxFactory:
module = import_module(mod_path)
cls = getattr(module, cls_name, None)
if not callable(cls):
raise TypeError(f"{cls_name} not found or not callable in {mod_path}")
return cast(PtxFactory, cls)
def _parse_backend_spec(spec: str) -> PtxFactory:
"""
Support:
- 'pkg.mod:Class'
- 'pkg.mod.Class'
"""
if ":" in spec:
left, cls_name = spec.split(":", 1)
return _load_from_dotted(left, cls_name)
if "." in spec:
mod_path, cls_name = spec.rsplit(".", 1)
return _load_from_dotted(mod_path, cls_name)
raise ValueError(
f"Unknown backend '{spec}'. "
f"Expected a registered name, 'pkg.mod:Class' or 'pkg.mod.Class'."
)
def get_backend(name: Optional[str] = None) -> ProtenixAPI:
"""
Priority:
1) explicit name
2) env variable PXDBENCH_BACKEND
3) default 'public'
"""
chosen = name or os.getenv("PXDBENCH_BACKEND") or "public"
if chosen in _FACTORIES:
return _FACTORIES[chosen]
return _parse_backend_spec(chosen)
|