Allow passing a base_schema to load_config()

RELEASE_NOTES.md

@@ -6,11 +6,12 @@
 
 ## Upgrading
 
-<!-- Here goes notes on how to upgrade from previous versions, including deprecations and what they should be replaced with -->
+- `frequenz.sdk.config.load_config()` doesn't accept classes decorated with `marshmallow_dataclass.dataclass` anymore. You should use the built-in `dataclasses.dataclass` directly instead, no other changes should be needed, the metadata in the `dataclass` fields will still be used.
 
 ## New Features
 
-<!-- Here goes the main new features and examples or instructions on how to use them -->
+
+- `frequenz.sdk.config.load_config()` can now use a base schema to customize even further how data is loaded.
 
 ## Bug Fixes
 

src/frequenz/sdk/config/_util.py

@@ -4,42 +4,71 @@
 """Utilities to deal with configuration."""
 
 from collections.abc import Mapping
-from typing import Any, TypeVar, cast
+from typing import Any, ClassVar, Protocol, TypeVar, cast
 
+from marshmallow import Schema
 from marshmallow_dataclass import class_schema
 
-T = TypeVar("T")
+
+# This is a hack that relies on identifying dataclasses by looking into an undocumented
+# property of dataclasses[1], so it might break in the future. Nevertheless, it seems to
+# be widely used in the community, for example `mypy` and `pyright` seem to rely on
+# it[2].
+#
+# [1]: https://github.com/python/mypy/issues/15974#issuecomment-1694781006
+# [2]: https://github.com/python/mypy/issues/15974#issuecomment-1694993493
+class Dataclass(Protocol):
+    """A protocol for dataclasses."""
+
+    __dataclass_fields__: ClassVar[dict[str, Any]]
+    """The fields of the dataclass."""
+
+
+DataclassT = TypeVar("DataclassT", bound=Dataclass)
 """Type variable for configuration classes."""
 
 
 def load_config(
-    cls: type[T],
+    cls: type[DataclassT],
     config: Mapping[str, Any],
     /,
+    base_schema: type[Schema] | None = None,
     **marshmallow_load_kwargs: Any,
-) -> T:
+) -> DataclassT:
     """Load a configuration from a dictionary into an instance of a configuration class.
 
     The configuration class is expected to be a [`dataclasses.dataclass`][], which is
     used to create a [`marshmallow.Schema`][] schema to validate the configuration
-    dictionary.
+    dictionary using [`marshmallow_dataclass.class_schema`][] (which in turn uses the
+    [`marshmallow.Schema.load`][] method to do the validation and deserialization).
 
-    To customize the schema derived from the configuration dataclass, you can use
-    [`marshmallow_dataclass.dataclass`][] to specify extra metadata.
+    To customize the schema derived from the configuration dataclass, you can use the
+    `metadata` key in [`dataclasses.field`][] to pass extra options to
+    [`marshmallow_dataclass`][] to be used during validation and deserialization.
 
     Additional arguments can be passed to [`marshmallow.Schema.load`][] using keyword
-    arguments.
+    arguments `marshmallow_load_kwargs`.
+
+    Note:
+        This method will raise [`marshmallow.ValidationError`][] if the configuration
+        dictionary is invalid and you have to have in mind all of the gotchas of
+        [`marshmallow`][] and [`marshmallow_dataclass`][] applies when using this
+        function.  It is recommended to carefully read the documentation of these
+        libraries.
 
     Args:
         cls: The configuration class.
         config: The configuration dictionary.
+        base_schema: An optional class to be used as a base schema for the configuration
+            class. This allow using custom fields for example. Will be passed to
+            [`marshmallow_dataclass.class_schema`][].
         **marshmallow_load_kwargs: Additional arguments to be passed to
             [`marshmallow.Schema.load`][].
 
     Returns:
         The loaded configuration as an instance of the configuration class.
     """
-    instance = class_schema(cls)().load(config, **marshmallow_load_kwargs)
+    instance = class_schema(cls, base_schema)().load(config, **marshmallow_load_kwargs)
     # We need to cast because `.load()` comes from marshmallow and doesn't know which
     # type is returned.
-    return cast(T, instance)
+    return cast(DataclassT, instance)

tests/config/test_util.py

@@ -7,7 +7,6 @@
 from typing import Any
 
 import marshmallow
-import marshmallow_dataclass
 import pytest
 from pytest_mock import MockerFixture
 
@@ -18,14 +17,6 @@
 class SimpleConfig:
     """A simple configuration class for testing."""
 
-    name: str
-    value: int
-
-
-@marshmallow_dataclass.dataclass
-class MmSimpleConfig:
-    """A simple configuration class for testing."""
-
     name: str = dataclasses.field(metadata={"validate": lambda s: s.startswith("test")})
     value: int
 
@@ -37,27 +28,36 @@ def test_load_config_dataclass() -> None:
     loaded_config = load_config(SimpleConfig, config)
     assert loaded_config == SimpleConfig(name="test", value=42)
 
-    config["name"] = "not test"
-    loaded_config = load_config(SimpleConfig, config)
-    assert loaded_config == SimpleConfig(name="not test", value=42)
-
-
-def test_load_config_marshmallow_dataclass() -> None:
-    """Test that load_config loads a configuration into a configuration class."""
-    config: dict[str, Any] = {"name": "test", "value": 42}
-    loaded_config = load_config(MmSimpleConfig, config)
-    assert loaded_config == MmSimpleConfig(name="test", value=42)
-
     config["name"] = "not test"
     with pytest.raises(marshmallow.ValidationError):
-        _ = load_config(MmSimpleConfig, config)
+        _ = load_config(SimpleConfig, config)
 
 
 def test_load_config_load_None() -> None:
     """Test that load_config raises ValidationError if the configuration is None."""
     config: dict[str, Any] = {}
     with pytest.raises(marshmallow.ValidationError):
-        _ = load_config(MmSimpleConfig, config.get("loggers", None))
+        _ = load_config(SimpleConfig, config.get("loggers", None))
+
+
+def test_load_config_with_base_schema() -> None:
+    """Test that load_config loads a configuration using a base schema."""
+
+    class _MyBaseSchema(marshmallow.Schema):
+        """A base schema for testing."""
+
+        class Meta:
+            """Meta options for the schema."""
+
+            unknown = marshmallow.EXCLUDE
+
+    config: dict[str, Any] = {"name": "test", "value": 42, "extra": "extra"}
+
+    loaded_config = load_config(SimpleConfig, config, base_schema=_MyBaseSchema)
+    assert loaded_config == SimpleConfig(name="test", value=42)
+
+    with pytest.raises(marshmallow.ValidationError):
+        _ = load_config(SimpleConfig, config)
 
 
 def test_load_config_type_hints(mocker: MockerFixture) -> None:
@@ -70,7 +70,7 @@ def test_load_config_type_hints(mocker: MockerFixture) -> None:
     config: dict[str, Any] = {}
 
     # We add the type hint to test that the return type (hint) is correct
-    _: MmSimpleConfig = load_config(MmSimpleConfig, config, marshmallow_arg=1)
+    _: SimpleConfig = load_config(SimpleConfig, config, marshmallow_arg=1)
     mock_class_schema.return_value.load.assert_called_once_with(
         config, marshmallow_arg=1
     )