""" Stage registry - Unified registration for all pipeline stages. Provides a single registry for sources, effects, displays, and cameras. Supports plugin discovery via entry points and explicit registration. """ from __future__ import annotations import importlib import importlib.metadata import logging from typing import TYPE_CHECKING, Any, TypeVar from sideline.pipeline.core import Stage if TYPE_CHECKING: from sideline.pipeline.core import Stage T = TypeVar("T") logger = logging.getLogger(__name__) class StageRegistry: """Unified registry for all pipeline stage types. Supports both explicit registration and automatic discovery via entry points. Plugins can be registered manually or discovered automatically. """ _categories: dict[str, dict[str, type[Any]]] = {} _discovered: bool = False _instances: dict[str, Stage] = {} _plugins_discovered: bool = False _plugin_modules: set[str] = set() # Track loaded plugin modules @classmethod def register(cls, category: str, stage_class: type[Any]) -> None: """Register a stage class in a category. Args: category: Category name (source, effect, display, camera) stage_class: Stage subclass to register """ if category not in cls._categories: cls._categories[category] = {} key = getattr(stage_class, "__name__", stage_class.__class__.__name__) cls._categories[category][key] = stage_class @classmethod def get(cls, category: str, name: str) -> type[Any] | None: """Get a stage class by category and name.""" return cls._categories.get(category, {}).get(name) @classmethod def list(cls, category: str) -> list[str]: """List all stage names in a category.""" return list(cls._categories.get(category, {}).keys()) @classmethod def list_categories(cls) -> list[str]: """List all registered categories.""" return list(cls._categories.keys()) @classmethod def create(cls, category: str, name: str, **kwargs) -> Stage | None: """Create a stage instance by category and name.""" stage_class = cls.get(category, name) if stage_class: return stage_class(**kwargs) return None @classmethod def create_instance(cls, stage: Stage | type[Stage], **kwargs) -> Stage: """Create an instance from a stage class or return as-is.""" if isinstance(stage, Stage): return stage if isinstance(stage, type) and issubclass(stage, Stage): return stage(**kwargs) raise TypeError(f"Expected Stage class or instance, got {type(stage)}") @classmethod def register_instance(cls, name: str, stage: Stage) -> None: """Register a stage instance by name.""" cls._instances[name] = stage @classmethod def get_instance(cls, name: str) -> Stage | None: """Get a registered stage instance by name.""" return cls._instances.get(name) @classmethod def register_plugin_module(cls, plugin_module: str) -> None: """Register stages from an external plugin module. The module should define a register_stages(registry) function. Args: plugin_module: Full module path (e.g., 'engine.plugins') """ if plugin_module in cls._plugin_modules: logger.debug(f"Plugin module {plugin_module} already loaded") return try: module = importlib.import_module(plugin_module) if hasattr(module, "register_stages"): module.register_stages(cls) cls._plugin_modules.add(plugin_module) logger.info(f"Registered stages from {plugin_module}") else: logger.warning( f"Plugin module {plugin_module} has no register_stages function" ) except ImportError as e: logger.warning(f"Failed to import plugin module {plugin_module}: {e}") # Backward compatibility alias register_plugin = register_plugin_module @classmethod def discover_plugins(cls) -> None: """Auto-discover and register plugins via entry points. Looks for 'sideline.stages' entry points in installed packages. Each entry point should point to a register_stages(registry) function. """ if cls._plugins_discovered: return try: # Discover entry points for sideline.stages # Python 3.12+ changed the entry_points() API try: entry_points = importlib.metadata.entry_points() if hasattr(entry_points, "get"): # Python < 3.12 stages_eps = entry_points.get("sideline.stages", []) else: # Python 3.12+ stages_eps = entry_points.select(group="sideline.stages") except Exception: # Fallback: try both approaches try: entry_points = importlib.metadata.entry_points() stages_eps = entry_points.get("sideline.stages", []) except Exception: stages_eps = [] for ep in stages_eps: try: register_func = ep.load() if callable(register_func): register_func(cls) logger.info(f"Discovered and registered plugin: {ep.name}") except Exception as e: logger.warning(f"Failed to load entry point {ep.name}: {e}") cls._plugins_discovered = True except Exception as e: logger.warning(f"Failed to discover plugins: {e}") @classmethod def get_discovered_modules(cls) -> set[str]: """Get set of plugin modules that have been loaded.""" return cls._plugin_modules.copy() def discover_stages() -> None: """Auto-discover and register all stage implementations. This function now only registers framework-level stages (displays, etc.). Application-specific stages should be registered via plugins. """ if StageRegistry._discovered: return # Register display stages (framework-level) _register_display_stages() # Discover plugins via entry points StageRegistry.discover_plugins() StageRegistry._discovered = True def _register_display_stages() -> None: """Register display backends as stages.""" try: from sideline.display import DisplayRegistry except ImportError: return DisplayRegistry.initialize() for backend_name in DisplayRegistry.list_backends(): factory = _DisplayStageFactory(backend_name) StageRegistry._categories.setdefault("display", {})[backend_name] = factory class _DisplayStageFactory: """Factory that creates DisplayStage instances for a specific backend.""" def __init__(self, backend_name: str): self._backend_name = backend_name def __call__(self): from sideline.display import DisplayRegistry from sideline.pipeline.adapters import DisplayStage display = DisplayRegistry.create(self._backend_name) if display is None: raise RuntimeError( f"Failed to create display backend: {self._backend_name}" ) return DisplayStage(display, name=self._backend_name) @property def __name__(self) -> str: return self._backend_name.capitalize() + "Stage" # Convenience functions def register_source(stage_class: type[Stage]) -> None: """Register a source stage.""" StageRegistry.register("source", stage_class) def register_effect(stage_class: type[Stage]) -> None: """Register an effect stage.""" StageRegistry.register("effect", stage_class) def register_display(stage_class: type[Stage]) -> None: """Register a display stage.""" StageRegistry.register("display", stage_class) def register_camera(stage_class: type[Stage]) -> None: """Register a camera stage.""" StageRegistry.register("camera", stage_class)