Looking for best design: Dynamically loading config and using configuration items as default arguments for functions

[sliedes]

At a high level, this is obviously more a design question than about omegaconf specifically, i.e. I don't see what omegaconf could do here to help, but I'd love to discuss it and ask for ideas about better ways to do what I want to do. I think I have a neat solution, with the problem that it's possibly too neat and magic, and perhaps I should do something different (I'll give you the code below). After all, I doubt my problem is uncommon. Even if there's no better approach, perhaps someone can get some good ideas from my current code.

So, I'd like to solicit your ideas or insight on what is the best way to solve it. I'm open to radically different solutions. Are there (arguably) better ways to solve this?

Problem

I want to use values loaded from config as default values of function arguments:

def foo(arg: int = config.default_arg) -> None: ...

Now this is obviously possible with omegaconf if we load the config at module import time, but for many reasons (e.g. testability) I'd actually very much like to avoid this.

I find it gets very unwieldy to everywhere do things like

def foo(arg: int | None = None) -> None:
    arg = config.default_arg if arg is None else arg
    ....

So, to recap, I would like to have both:

1. Avoid loading config at import time, so that for example my tests do not depend on a config
2. Use a value from the config for a (module level) function default parameter.

In Python, default values are evaluated at import time, so this obviously is not quite possible exactly as stated.

My current solution

This feels too magic, but kind of works. What I like is that it avoids the repetitive arg = ... if arg is None else ... mess. It allows me to achieve what I want using syntax like this:

@uses_config
def foo(arg: int = FromConfig("default_arg")) -> None: ...

The wrapper uses_config replaces FromConfig default values with values from the config at first call time (not wrapper creation time), at a point where the configuration should already be loaded.

See Appendix for the code implementing this.

Discussion

There are certainly things that I don't like here, starting from how magic this is; the fact that I lose some type safety (though I think the wrapper could be modified to compare the type annotations); and that I need to remember to use @uses_config.

I think one alternative approach, with its own downsides, would be to get rid of the global config object entirely (globals kind of suck, not least for testability), make load return a _Config and then pass that to everything. I don't immediately see how this would help me get rid of the magic, but it would allow us to get rid of the delicate "do not call this at import time" pitfall.

I also considered using typing.Annotations to mark the arguments that should have default values taken from the config, but it feels like that doesn't really make it easier in any way.

My initial approach actually involved creating a _Config object named fromConfig with the same hierarchy as a real config, but with all the leaf values replaced by objects that told the path to the config. Then you could say

@uses_config
def foo(arg: int = fromConfig.backend.port) -> None: ...

and even have type safety. Unfortunately, creating such a class to mirror an attrs class (with some validation) turned out to be a bit tricky, and of course it still won't work in expressions like fromConfig.backend.port + 1. Lists and dictionaries of items in the config seem to be tricky, too.

Appendix: Implementing code

import contextlib
import functools
import inspect
from pathlib import Path
from typing import IO, Any, Callable, TypeVar, cast

import attrs
from loguru import logger
from omegaconf import OmegaConf

CONFIG_PATH = Path("config.yaml")


@attrs.define
class DocumentSync:
    """An example config class."""

    data_gen_path: Path
    source_docs_path: Path
    parsed_docs_path: Path


@attrs.define
class Backend:
    """An example config class."""

    listen_port: int


@attrs.define
class _Config:
    """An example top level config class."""

    backend: Backend
    document_sync: DocumentSync


class ConfigurationNotLoadedError(Exception):
    """Raised when the configuration is not loaded yet when trying to access it."""


class _ConfigProxy:
    """
    A proxy for the config object before the config is loaded.

    At attribute access time, checks if the config has since been loaded (can happen if this is "imported from" before
    the config is loaded) and forwards the attribute access to the config object.
    """

    def __getattribute__(self, name: str) -> Any:
        global config

        if name.startswith("__"):
            return super().__getattribute__(name)

        # This can happen if this has been "imported from" before the config was loaded
        if isinstance(config, _Config):
            return getattr(config, name)

        logger.error("Configuration not loaded yet")
        raise ConfigurationNotLoadedError("Configuration not loaded yet")


def _is_import_in_progress() -> bool:
    return any("importlib._bootstrap" in frame_info.filename for frame_info in inspect.stack())


# The config object to be used to access the config.
config: _Config = cast(_Config, _ConfigProxy())


def load(
    *,
    reload: bool = False,
    config_file: Path | IO[Any] = CONFIG_PATH,
    _allow_module_level: bool = False,  # for testing only
) -> None:
    """
    (Re)load the configuration.

    Do not call from module level (at import time).
    """

    global config

    if isinstance(config, _Config) and not reload:
        return

    if not _allow_module_level and _is_import_in_progress():
        raise ConfigurationNotLoadedError("Configuration loaded during import")

    config = cast(
        _Config, OmegaConf.to_object(OmegaConf.merge(OmegaConf.structured(_Config), OmegaConf.load(config_file)))
    )


@attrs.frozen
class _FromConfig:
    """A marker for a function argument that should be replaced by a value from the config."""

    key: str


FromConfig: Callable[[str], Any] = _FromConfig


@functools.cache
def _get_config_item(path: str) -> Any:
    cfg = config
    for part in path.split("."):
        cfg = getattr(cfg, part)
    return cfg


_T = TypeVar("_T")


def _check_config_key(path: str) -> None:
    cls: type[Any] = _Config
    for part in path.split("."):
        with contextlib.suppress(attrs.exceptions.NotAnAttrsClassError):
            cls = getattr(attrs.fields(cls), part).type


def uses_config(func: Callable[..., _T]) -> Callable[..., _T]:
    """Replace uses of FromConfig by values from config."""

    config_args: dict[str, Any] | None = None

    # Validate the config keys at wrapper creation time
    for arg_value in inspect.signature(func).parameters.values():
        if isinstance(arg_value.default, _FromConfig):
            _check_config_key(arg_value.default.key)

    def init() -> None:
        # Initialize the config args at first call time
        nonlocal config_args

        if _is_import_in_progress():
            raise ConfigurationNotLoadedError("Function with @uses_config decorator called during import")

        config_args = {}
        for arg_name, arg_value in inspect.signature(func).parameters.items():
            if isinstance(arg_value.default, _FromConfig):
                config_args[arg_name] = _get_config_item(arg_value.default.key)

    @functools.wraps(func)
    def wrapper(*args: tuple[Any, ...], **kwargs: dict[str, Any]) -> _T:
        if config_args is None:
            init()
        assert config_args is not None
        kwargs = {**config_args, **kwargs}
        return func(*args, **kwargs)

    return wrapper