from typing import Dict from pydantic import BaseModel class Registry(BaseModel): """Registry for storing and building classes.""" name: str entries: Dict = {} def register(self, key: str): def decorator(class_builder): self.entries[key] = class_builder return class_builder return decorator def build(self, type: str, **kwargs): if type not in self.entries: raise ValueError( f'{type} is not registered. Please register with the .register("{type}") method provided in {self.name} registry' ) return self.entries[type](**kwargs) def get_all_entries(self): return self.entries