SingingSDS / data_handlers /__init__.py
jhansss's picture
refactor init
91394e0
raw
history blame contribute delete
807 Bytes
import importlib
import pkgutil
from pathlib import Path
from .base import MelodyDatasetHandler
_registry = {}
for _, module_name, _ in pkgutil.iter_modules([str(Path(__file__).parent)]):
if module_name in ("__init__", "base"):
continue
module = importlib.import_module(f"{__name__}.{module_name}")
for attr_name in dir(module):
attr = getattr(module, attr_name)
if (
isinstance(attr, type)
and issubclass(attr, MelodyDatasetHandler)
and attr is not MelodyDatasetHandler
):
_registry[attr.name] = attr # 注册 class 本身
def get_melody_handler(name: str) -> type[MelodyDatasetHandler]:
if name not in _registry:
raise ValueError(f"Melody source '{name}' not found")
return _registry[name]