Add method to registry to get subclasses of a class.

pyaml/validation/registry.py

@@ -59,6 +59,8 @@ class SchemaRegistry:
         Return a view of registered class paths.
     values()
         Return a view of registered schemas.
+    subclasses_of(schema, include_base=False)
+        Return registered schemas derived from a schema.
     update(class_path, schema)
         Replace the schema registered for a class path.
     """
@@ -315,6 +317,44 @@ def values(
         """
         return self._schemas.values()
 
+    def subclasses_of(
+        self,
+        schema: type[ConfigurationSchema],
+        *,
+        include_base: bool = False,
+    ) -> dict[str, type[ConfigurationSchema]]:
+        """Return registered schemas that are subclasses of ``schema``.
+
+        Registered virtual subclasses are included. By default, ``schema`` itself
+        is excluded; set ``include_base`` to ``True`` to include it when registered.
+
+        Parameters
+        ----------
+        schema : type[ConfigurationSchema]
+            Base schema to search for.
+        include_base : bool, optional
+            Include an entry whose schema is exactly ``schema``.
+
+        Returns
+        -------
+        dict[str, type[ConfigurationSchema]]
+            Mapping of registered class paths to matching schema classes.
+
+        Raises
+        ------
+        TypeError
+            If ``schema`` is not a ``ConfigurationSchema`` subclass.
+        """
+        if not isinstance(schema, type) or not issubclass(schema, ConfigurationSchema):
+            raise TypeError(f"{schema!r} must inherit from ConfigurationSchema.")
+
+        return {
+            class_path: registered_schema
+            for class_path, registered_schema in self._schemas.items()
+            if (include_base or registered_schema is not schema)
+            and (issubclass(registered_schema, schema) or registered_schema.is_virtual_subclass_of(schema))
+        }
+
     def __len__(
         self,
     ) -> int:

tests/validation/test_registry.py

@@ -178,6 +178,27 @@ def test_getitem_raises_clean_keyerror_for_missing_schema(registry: SchemaRegist
         _ = registry["pkg.module.Class"]
 
 
+def test_subclasses_of_returns_registered_concrete_schemas(registry: SchemaRegistry):
+    class ConcreteSchema(DummySchema):
+        pass
+
+    registry.register("pkg.module.Base", DummySchema)
+    registry.register("pkg.module.Concrete", ConcreteSchema)
+    registry.register("pkg.module.Other", OtherSchema)
+
+    assert registry.subclasses_of(DummySchema) == {
+        "pkg.module.Concrete": ConcreteSchema,
+    }
+
+
+def test_subclasses_of_can_include_base_schema(registry: SchemaRegistry):
+    registry.register("pkg.module.Base", DummySchema)
+
+    assert registry.subclasses_of(DummySchema, include_base=True) == {
+        "pkg.module.Base": DummySchema,
+    }
+
+
 def test_get_returns_registered_schema(registry: SchemaRegistry):
     registry.register("pkg.module.Class", DummySchema)