Make load_config() accept a base_schema

This is necessary to use custom fields when loading configurations.

Signed-off-by: Leandro Lucarella <[email protected]>

Author: llucax

RELEASE_NOTES.md

@@ -10,7 +10,8 @@
 
 ## 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

@@ -6,6 +6,7 @@
 from collections.abc import Mapping
 from typing import Any, TypeVar, cast
 
+from marshmallow import Schema
 from marshmallow_dataclass import class_schema
 
 T = TypeVar("T")
@@ -16,6 +17,7 @@ def load_config(
     cls: type[T],
     config: Mapping[str, Any],
     /,
+    base_schema: type[Schema] | None = None,
     **marshmallow_load_kwargs: Any,
 ) -> T:
     """Load a configuration from a dictionary into an instance of a configuration class.
@@ -42,13 +44,16 @@ def load_config(
     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)

tests/config/test_util.py

@@ -37,7 +37,27 @@ 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: