quick utils enhancements

Author: markurtz

src/guidellm/utils/pydantic_utils.py

@@ -28,6 +28,7 @@
 
 
 BaseModelT = TypeVar("BaseModelT", bound=BaseModel)
+RegisterClassT = TypeVar("RegisterClassT", bound=type[BaseModelT])
 SuccessfulT = TypeVar("SuccessfulT")
 ErroredT = TypeVar("ErroredT")
 IncompleteT = TypeVar("IncompleteT")
@@ -130,7 +131,7 @@ class StatusBreakdown(BaseModel, Generic[SuccessfulT, ErroredT, IncompleteT, Tot
 
     Example:
     ::
-        from guidellm.utils.pydantic_utils import StatusBreakdown
+        from guidellm.utils import StatusBreakdown
 
         # Define a breakdown for request counts
         breakdown = StatusBreakdown[int, int, int, int](
@@ -172,7 +173,7 @@ class PydanticClassRegistryMixin(
 
     Example:
     ::
-        from guidellm.utils.pydantic_utils import PydanticClassRegistryMixin
+        from speculators.utils import PydanticClassRegistryMixin
 
         class BaseConfig(PydanticClassRegistryMixin["BaseConfig"]):
             schema_discriminator: ClassVar[str] = "config_type"
@@ -200,8 +201,8 @@ class DatabaseConfig(BaseConfig):
 
     @classmethod
     def register_decorator(
-        cls, clazz: type[BaseModelT], name: str | list[str] | None = None
-    ) -> type[BaseModelT]:
+        cls, clazz: RegisterClassT, name: str | list[str] | None = None
+    ) -> RegisterClassT:
         """
         Register a Pydantic model class with type validation and schema reload.
 
@@ -300,3 +301,25 @@ def auto_populate_registry(cls) -> bool:
         cls.reload_schema()
 
         return populated
+
+    @classmethod
+    def registered_classes(cls) -> tuple[type[Any], ...]:
+        """
+        Get all registered pydantic classes from the registry.
+
+        Automatically triggers auto-discovery if registry_auto_discovery is enabled
+        to ensure all available implementations are included.
+
+        :return: Tuple of all registered classes including auto-discovered ones
+        :raises ValueError: If called before any objects have been registered
+        """
+        if cls.registry_auto_discovery:
+            cls.auto_populate_registry()
+
+        if cls.registry is None:
+            raise ValueError(
+                "ClassRegistryMixin.registered_classes() must be called after "
+                "registering classes with ClassRegistryMixin.register()."
+            )
+
+        return tuple(cls.registry.values())

src/guidellm/utils/registry.py

@@ -14,23 +14,23 @@
 
 from guidellm.utils.auto_importer import AutoImporterMixin
 
-__all__ = ["RegistryMixin", "RegistryObjT"]
+__all__ = ["RegisterT", "RegistryMixin", "RegistryObjT"]
 
 
 RegistryObjT = TypeVar("RegistryObjT", bound=Any)
-"""
-Generic type variable for objects managed by the registry system.
-"""
+"""Generic type variable for objects managed by the registry system."""
+RegisterT = TypeVar("RegisterT", bound=RegistryObjT)
+"""Generic type variable for the args and return values within the registry."""
 
 
 class RegistryMixin(Generic[RegistryObjT], AutoImporterMixin):
     """
     Generic mixin for creating object registries with optional auto-discovery.
 
-    Enables classes to maintain separate registries of objects that can be
-    dynamically discovered and instantiated through decorators and module imports.
-    Supports both manual registration via decorators and automatic discovery
-    through package scanning for extensible plugin architectures.
+    Enables classes to maintain separate registries of objects that can be dynamically
+    discovered and instantiated through decorators and module imports. Supports both
+    manual registration via decorators and automatic discovery through package scanning
+    for extensible plugin architectures.
 
     Example:
     ::
@@ -69,14 +69,14 @@ class TokenProposal(RegistryMixin):
     @classmethod
     def register(
         cls, name: str | list[str] | None = None
-    ) -> Callable[[RegistryObjT], RegistryObjT]:
+    ) -> Callable[[RegisterT], RegisterT]:
         """
-        Decorator that registers an object with the registry.
+        Decorator for registering objects with the registry.
 
         :param name: Optional name(s) to register the object under.
-            If None, the object name is used as the registry key.
-        :return: A decorator function that registers the decorated object.
-        :raises ValueError: If name is provided but is not a string or list of strings.
+            If None, uses the object's __name__ attribute
+        :return: Decorator function that registers the decorated object
+        :raises ValueError: If name is not a string, list of strings, or None
         """
         if name is not None and not isinstance(name, (str, list)):
             raise ValueError(
@@ -88,19 +88,19 @@ def register(
 
     @classmethod
     def register_decorator(
-        cls, obj: RegistryObjT, name: str | list[str] | None = None
-    ) -> RegistryObjT:
+        cls, obj: RegisterT, name: str | list[str] | None = None
+    ) -> RegisterT:
         """
-        Direct decorator that registers an object with the registry.
+        Register an object directly with the registry.
 
-        :param obj: The object to register.
+        :param obj: The object to register
         :param name: Optional name(s) to register the object under.
-            If None, the object name is used as the registry key.
-        :return: The registered object.
-        :raises ValueError: If the object is already registered or if name is invalid.
+            If None, uses the object's __name__ attribute
+        :return: The registered object
+        :raises ValueError: If the object is already registered or name is invalid
         """
 
-        if not name:
+        if name is None:
             name = obj.__name__
         elif not isinstance(name, (str, list)):
             raise ValueError(
@@ -127,20 +127,20 @@ def register_decorator(
                     "registered."
                 )
 
-            cls.registry[register_name.lower()] = obj
+            cls.registry[register_name] = obj
 
         return obj
 
     @classmethod
     def auto_populate_registry(cls) -> bool:
         """
-        Import and register all modules from the specified auto_package.
+        Import and register all modules from the auto_package.
 
         Automatically called by registered_objects when registry_auto_discovery is True
-        to ensure all available implementations are discovered before returning results.
+        to ensure all available implementations are discovered.
 
-        :return: True if the registry was populated, False if already populated.
-        :raises ValueError: If called when registry_auto_discovery is False.
+        :return: True if registry was populated, False if already populated
+        :raises ValueError: If called when registry_auto_discovery is False
         """
         if not cls.registry_auto_discovery:
             raise ValueError(
@@ -165,8 +165,8 @@ def registered_objects(cls) -> tuple[RegistryObjT, ...]:
         Automatically triggers auto-discovery if registry_auto_discovery is enabled
         to ensure all available implementations are included.
 
-        :return: Tuple of all registered objects including auto-discovered ones.
-        :raises ValueError: If called before any objects have been registered.
+        :return: Tuple of all registered objects including auto-discovered ones
+        :raises ValueError: If called before any objects have been registered
         """
         if cls.registry_auto_discovery:
             cls.auto_populate_registry()
@@ -183,24 +183,33 @@ def registered_objects(cls) -> tuple[RegistryObjT, ...]:
     def is_registered(cls, name: str) -> bool:
         """
         Check if an object is registered under the given name.
+        It matches first by exact name, then by str.lower().
 
         :param name: The name to check for registration.
         :return: True if the object is registered, False otherwise.
         """
         if cls.registry is None:
             return False
 
-        return name.lower() in cls.registry
+        return name in cls.registry or name.lower() in [
+            key.lower() for key in cls.registry
+        ]
 
     @classmethod
     def get_registered_object(cls, name: str) -> RegistryObjT | None:
         """
-        Get a registered object by its name.
+        Get a registered object by its name. It matches first by exact name,
+        then by str.lower().
 
         :param name: The name of the registered object.
         :return: The registered object if found, None otherwise.
         """
         if cls.registry is None:
             return None
 
-        return cls.registry.get(name.lower())
+        if name in cls.registry:
+            return cls.registry[name]
+
+        lower_key_map = {key.lower(): key for key in cls.registry}
+
+        return cls.registry.get(lower_key_map.get(name.lower()))

tests/unit/utils/test_pydantic_utils.py

@@ -10,7 +10,7 @@
 import pytest
 from pydantic import BaseModel, Field, ValidationError
 
-from guidellm.utils.pydantic_utils import (
+from guidellm.utils import (
     PydanticClassRegistryMixin,
     ReloadableBaseModel,
     StandardBaseDict,
@@ -459,6 +459,7 @@ def test_class_signatures(self):
         assert hasattr(PydanticClassRegistryMixin, "__get_pydantic_core_schema__")
         assert hasattr(PydanticClassRegistryMixin, "__pydantic_generate_base_schema__")
         assert hasattr(PydanticClassRegistryMixin, "auto_populate_registry")
+        assert hasattr(PydanticClassRegistryMixin, "registered_classes")
 
     @pytest.mark.smoke
     def test_initialization(self, valid_instances):
@@ -547,8 +548,8 @@ class TestSubModel(TestBaseModel):
             value: str
 
         assert TestBaseModel.registry is not None  # type: ignore[misc]
-        assert "testsubmodel" in TestBaseModel.registry  # type: ignore[misc]
-        assert TestBaseModel.registry["testsubmodel"] is TestSubModel  # type: ignore[misc]
+        assert "TestSubModel" in TestBaseModel.registry  # type: ignore[misc]
+        assert TestBaseModel.registry["TestSubModel"] is TestSubModel  # type: ignore[misc]
 
     @pytest.mark.sanity
     def test_register_decorator_with_name(self):
@@ -621,6 +622,87 @@ def __pydantic_schema_base_type__(cls) -> type[TestBaseModel]:
             assert result is True
             mock_reload.assert_called_once()
 
+    @pytest.mark.smoke
+    def test_registered_classes(self):
+        """Test PydanticClassRegistryMixin.registered_classes method."""
+
+        class TestBaseModel(PydanticClassRegistryMixin):
+            schema_discriminator: ClassVar[str] = "test_type"
+            test_type: str
+            registry_auto_discovery: ClassVar[bool] = False
+
+            @classmethod
+            def __pydantic_schema_base_type__(cls) -> type[TestBaseModel]:
+                if cls.__name__ == "TestBaseModel":
+                    return cls
+                return TestBaseModel
+
+        @TestBaseModel.register("test_sub_a")
+        class TestSubModelA(TestBaseModel):
+            test_type: str = "test_sub_a"
+            value_a: str
+
+        @TestBaseModel.register("test_sub_b")
+        class TestSubModelB(TestBaseModel):
+            test_type: str = "test_sub_b"
+            value_b: int
+
+        # Test normal case with registered classes
+        registered = TestBaseModel.registered_classes()
+        assert isinstance(registered, tuple)
+        assert len(registered) == 2
+        assert TestSubModelA in registered
+        assert TestSubModelB in registered
+
+    @pytest.mark.sanity
+    def test_registered_classes_with_auto_discovery(self):
+        """Test PydanticClassRegistryMixin.registered_classes with auto discovery."""
+
+        class TestBaseModel(PydanticClassRegistryMixin):
+            schema_discriminator: ClassVar[str] = "test_type"
+            test_type: str
+            registry_auto_discovery: ClassVar[bool] = True
+
+            @classmethod
+            def __pydantic_schema_base_type__(cls) -> type[TestBaseModel]:
+                if cls.__name__ == "TestBaseModel":
+                    return cls
+                return TestBaseModel
+
+        with mock.patch.object(
+            TestBaseModel, "auto_populate_registry"
+        ) as mock_auto_populate:
+            # Mock the registry to simulate registered classes
+            TestBaseModel.registry = {"test_class": type("TestClass", (), {})}
+            mock_auto_populate.return_value = False
+
+            registered = TestBaseModel.registered_classes()
+            mock_auto_populate.assert_called_once()
+            assert isinstance(registered, tuple)
+            assert len(registered) == 1
+
+    @pytest.mark.sanity
+    def test_registered_classes_no_registry(self):
+        """Test PydanticClassRegistryMixin.registered_classes with no registry."""
+
+        class TestBaseModel(PydanticClassRegistryMixin):
+            schema_discriminator: ClassVar[str] = "test_type"
+            test_type: str
+
+            @classmethod
+            def __pydantic_schema_base_type__(cls) -> type[TestBaseModel]:
+                if cls.__name__ == "TestBaseModel":
+                    return cls
+                return TestBaseModel
+
+        # Ensure registry is None
+        TestBaseModel.registry = None
+
+        with pytest.raises(ValueError) as exc_info:
+            TestBaseModel.registered_classes()
+
+        assert "must be called after registering classes" in str(exc_info.value)
+
     @pytest.mark.sanity
     def test_marshalling(self, valid_instances):
         """Test PydanticClassRegistryMixin serialization and deserialization."""
@@ -707,4 +789,4 @@ class ContainerModel(BaseModel):
         assert isinstance(recreated.model, TestSubModelA)
         assert len(recreated.models) == 2
         assert isinstance(recreated.models[0], TestSubModelA)
-        assert isinstance(recreated.models[1], TestSubModelB)
+        assert isinstance(recreated.models[1], TestSubModelB)