dotenv utility wrapper

env.py

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Any, Generic, Optional, TypeVar, cast, overload

from dotenv import dotenv_values  # type: ignore
from typing_extensions import LiteralString

KI = TypeVar("KI", bound=LiteralString)
KO = TypeVar("KO", bound=LiteralString)
KIB = TypeVar("KIB", bound=LiteralString)
KOB = TypeVar("KOB", bound=LiteralString)

# fmt: off
@overload
def env(key: KI, *, default: Optional[str] = None) -> EnvSpec[KI, KI]: ...
@overload
def env(key: KI, rename: None = None, default: Optional[str] = None) -> EnvSpec[KI, KI]: ...
@overload
def env(key: KI, rename: KO, default: Optional[str] = None) -> EnvSpec[KI, KO]: ...
# fmt: on


def env(key: Any, rename: Any = None, default: Optional[str] = None):
    """
    Small utility to get environment variables with type checking and fail early on missing.

    Provides optional renaming and optional default value if not present in env.

    If keys are missing, it raises with all the missing keys instead of just the first one.

    Usage:
    >>> spec = env("FOO") & env("BAR") & env("BAZ", "baz_renamed") & env("QUX", default="qux_default")
    >>> spec.load('../.env.local')
    Exception: missing env keys ['FOO', 'BAR', 'BAZ']

    >>> spec.parse({"FOO": "42", "BAR": "69", "BAZ": "420"})
    {'FOO': '42', 'BAR': '69', 'baz_renamed': '420', 'QUX': 'qux_default'}
    """
    return EnvSpec.create(key, rename or key, default)


@dataclass(frozen=True)
class EnvSpec(Generic[KI, KO]):
    renames: dict[KI, KO]
    defaults: dict[KI, Optional[str]]

    @classmethod
    def create(cls, kin: KI, kout: KO, default: Optional[str] = None) -> EnvSpec[KI, KO]:
        return cls(renames={kin: kout}, defaults={kin: default})

    def load(self, dotenv_path: str | Path) -> dict[KO, str]:
        return self.parse(dotenv_values(str(dotenv_path)))

    def parse(self, env: dict[Any, Any]) -> dict[KO, str]:
        result: dict[KO, str] = {}
        missing: list[str] = []

        for kin, kout in self.renames.items():
            if env.get(kin, None) or self.defaults[kin] is not None:
                result[kout] = cast(str, env.get(kin, self.defaults[kin]))
            else:
                missing.append(kin)

        assert len(missing) == 0, f"missing env keys {missing}"
        return result

    def __and__(self, other: EnvSpec[KIB, KOB]) -> EnvSpec[KI | KIB, KO | KOB]:
        """self & other"""
        return EnvSpec({**self.renames, **other.renames}, {**self.defaults, **other.defaults})