src/open_storyline/utils/register.py

# core/registry.py
import pkgutil
import importlib
from typing import Optional
import os

class Registry:
    def __init__(self):
        self._items = {}


    def register(self, name: Optional[str] = None, override: bool = False):
        """
        Class decorator to register a class in the registry.
        
        Args:
            name (str, optional): Custom registration name. Defaults to module_name.ClassName.
            override (bool): If True, will replace an existing class with the same name. Defaults to False.
        """
        def decorator(cls):
            reg_name = name or f"{cls.__name__}"
            if reg_name in self._items:
                if override:
                    print(f"[Registry] {reg_name} already registered, override=True -> replacing")
                else:
                    raise KeyError(f"[Registry] {reg_name} already registered, override=False")
            
            self._items[reg_name] = cls
            print(f"[Registry] Registered: {reg_name}")
            return cls
        return decorator


    def get(self, name: str, default=None):
        """Get a registered class by name. Returns `default` if not found."""
        return self._items.get(name, default)


    def list(self):
        """Return a list of all registered names."""
        return list(self._items.keys())


    def __len__(self):
        return len(self._items)


    def clear(self):
        """Clear all registered classes."""
        self._items.clear()


    def scan_package(self, package_name: str):
        """
        Scan a Python package and its subpackages, import modules to trigger @REGISTRY.register().
        
        Args:
            package_name (str): Name of the package, e.g., "nodes"
        """
        package = importlib.import_module(package_name)
        if not hasattr(package, "__path__"):
            # Not a package, skip scanning
            print(f"[Registry] {package_name} is not a package, skipping scan")
            return


        for finder, modname, ispkg in pkgutil.walk_packages(package.__path__, package.__name__ + "."):
            importlib.import_module(modname)
            print(f"[Registry] Scanned module: {modname}")


# Global registry instance
NODE_REGISTRY = Registry()