A simple way to load a TOML config with CLI arguments. For when you don't want to pull in hydra.

toml_w_overrides.py

"""Load TOML configuration files with optional dotted-path overrides.

The module reads a TOML file with :mod:`tomllib`, then applies zero or more
``key=value`` overrides on top of the parsed data.

Override keys use dotted paths to address nested tables:

.. code-block:: python

    >>> from pathlib import Path
    >>> path = Path("train.toml")
    >>> _ = path.write_text(
    ...     '[model]\\nname = "resnet18"\\n[optimizer]\\nlr = 0.001\\n',
    ...     encoding="utf-8",
    ... )
    >>> load(path, ['optimizer.lr=0.01', 'model.name="vit_b16"'])
    {'model': {'name': 'vit_b16'}, 'optimizer': {'lr': 0.01}}

Override values are parsed as TOML literals when possible, so numbers, booleans,
lists, and quoted strings keep their natural types:

.. code-block:: python

    >>> load(path, ["seed=17", "use_amp=true", "layers=[256, 128, 64]"])
    {
    ...     'model': {'name': 'resnet18'},
    ...     'optimizer': {'lr': 0.001},
    ...     'seed': 17,
    ...     'use_amp': True,
    ...     'layers': [256, 128, 64],
    ... }

If a value is not valid TOML syntax, it is kept as a plain string. This makes
shell-friendly arguments such as ``env=prod`` work without extra quoting.
"""

from __future__ import annotations

import sys
import tomllib
from pathlib import Path
from typing import Any, Sequence

Config = dict[str, Any]


def _parse_value(raw: str) -> Any:
    """Parse a TOML literal, or keep the original string if unquoted."""
    try:
        return tomllib.loads(f"value = {raw}")["value"]
    except tomllib.TOMLDecodeError:
        return raw


def _insert_path(target: Config, path: str, value: Any) -> None:
    keys = path.split(".")
    if not path or any(not key for key in keys):
        raise ValueError(
            "override keys must be non-empty and cannot contain empty path segments"
        )

    current = target
    for key in keys[:-1]:
        child = current.get(key)
        if not isinstance(child, dict):
            child = {}
            current[key] = child
        current = child
    current[keys[-1]] = value


def _parse_overrides(args: Sequence[str]) -> Config:
    overrides: Config = {}
    for arg in args:
        path, sep, raw = arg.partition("=")
        if not sep:
            continue
        try:
            _insert_path(overrides, path, _parse_value(raw))
        except ValueError as exc:
            raise ValueError(f"invalid override {arg!r}: {exc}") from exc
    return overrides


def _merge_dicts(base: Config, patch: Config) -> Config:
    merged = base.copy()
    for key, value in patch.items():
        if isinstance(value, dict) and isinstance(merged.get(key), dict):
            merged[key] = _merge_dicts(merged[key], value)
        else:
            merged[key] = value
    return merged


def load(
    path: str | Path,
    argv: Sequence[str] | None = None,
    *,
    parse_cli: bool = False,
) -> Config:
    """Load a TOML file and optionally merge dotted ``key=value`` overrides.

    Parameters
    ----------
    path:
        Path to the TOML file to read.
    argv:
        Override arguments in ``key=value`` form. Keys may use dotted paths such
        as ``database.pool.size=8``. If omitted, no overrides are applied unless
        ``parse_cli`` is true.
    parse_cli:
        When true and ``argv`` is not provided, overrides are read from
        ``sys.argv[1:]``.

    Returns
    -------
    dict[str, Any]
        The parsed TOML document with overrides merged in.

    Notes
    -----
    Values are parsed as TOML literals when possible. For example,
    ``steps=1000`` becomes an integer and ``use_amp=true`` becomes a boolean.
    A value that is not valid TOML syntax is left as a plain string, so
    ``device=cuda`` results in ``"cuda"``.
    """
    with Path(path).open("rb") as fh:
        config: Config = tomllib.load(fh)

    args = list(argv) if argv is not None else (sys.argv[1:] if parse_cli else [])
    if not args:
        return config

    return _merge_dicts(config, _parse_overrides(args))