pixel_cursor/ops.py

"""Op dispatch registry.

The cursor core defines the op contract; backends (in studio/backends/) register
concrete implementations via register_backend. Keeps pixel_cursor/ free of model
imports — torch/diffusers/MotionClone never enter the core package.
"""
from __future__ import annotations
from typing import Callable, Any


_REGISTRY: dict[tuple[str, str], Callable] = {}


def register_backend(op_name: str, backend_name: str, fn: Callable) -> None:
    _REGISTRY[(op_name, backend_name)] = fn


def list_backends(op_name: str | None = None) -> list[tuple[str, str]]:
    if op_name is None:
        return sorted(_REGISTRY.keys())
    return sorted(k for k in _REGISTRY if k[0] == op_name)


def _dispatch(op_name: str, backend: str, cursor, *args, **kwargs) -> Any:
    if backend == "auto":
        for (op, _name), fn in _REGISTRY.items():
            if op == op_name:
                return fn(cursor, *args, **kwargs)
        raise NotImplementedError(
            f"No backend registered for op {op_name!r}. "
            f"Register one via pixel_cursor.ops.register_backend()."
        )
    fn = _REGISTRY.get((op_name, backend))
    if fn is None:
        available = sorted({n for (o, n) in _REGISTRY if o == op_name})
        raise NotImplementedError(
            f"No backend {backend!r} for op {op_name!r}; available: {available}"
        )
    return fn(cursor, *args, **kwargs)