plugin.py

"""Adopt environment section in pytest configuration files."""

from __future__ import annotations

import os
import sys
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any

import pytest
from dotenv import dotenv_values

if TYPE_CHECKING:
    from collections.abc import Generator, Iterator
    from pathlib import Path

_env_actions_key = pytest.StashKey[list[str]]()

if sys.version_info >= (3, 11):  # pragma: >=3.11 cover
    import tomllib
else:  # pragma: <3.11 cover
    import tomli as tomllib


def pytest_addoption(parser: pytest.Parser) -> None:
    """Add section to configuration files."""
    help_msg = "a line separated list of environment variables of the form (FLAG:)NAME=VALUE"
    parser.addini("env", type="linelist", help=help_msg, default=[])
    parser.addini("env_files", type="linelist", help="a line separated list of .env files to load", default=[])
    parser.addini(
        "env_files_skip_if_set",
        type="bool",
        help="only set .env file variables when not already defined",
        default=False,
    )
    parser.addoption(
        "--envfile",
        action="store",
        dest="envfile",
        default=None,
        help="path to .env file to load (prefix with + to extend config files, otherwise replaces them)",
    )
    parser.addoption(
        "--pytest-env-verbose",
        action="store_true",
        dest="pytest_env_verbose",
        default=False,
        help="print environment variable assignments made by pytest-env",
    )


@dataclass
class Entry:
    """Configuration entries."""

    key: str
    value: str
    transform: bool
    skip_if_set: bool
    unset: bool = False


@pytest.hookimpl(tryfirst=True)
def pytest_load_initial_conftests(
    args: list[str],  # noqa: ARG001
    early_config: pytest.Config,
    parser: pytest.Parser,  # noqa: ARG001
) -> None:
    """Load environment variables from configuration files."""
    verbose = getattr(early_config.known_args_namespace, "pytest_env_verbose", False)
    actions: list[tuple[str, str, str, str]] = []

    env_files_list: list[str] = []
    env_files_skip_if_set: bool | None = None
    if toml_config := _find_toml_config(early_config):
        env_files_list, _, env_files_skip_if_set = _load_toml_config(toml_config)

    if env_files_skip_if_set is None:
        env_files_skip_if_set = bool(early_config.getini("env_files_skip_if_set"))

    _apply_env_files(
        early_config,
        env_files_list,
        actions if verbose else None,
        skip_if_set=env_files_skip_if_set,
    )
    _apply_entries(early_config, actions if verbose else None)

    if verbose and actions:
        early_config.stash[_env_actions_key] = _format_actions(actions)


def _apply_env_files(
    early_config: pytest.Config,
    env_files_list: list[str],
    actions: list[tuple[str, str, str, str]] | None,
    *,
    skip_if_set: bool = False,
) -> None:
    preexisting = dict(os.environ) if skip_if_set else {}
    for env_file in _load_env_files(early_config, env_files_list):
        for key, value in dotenv_values(env_file).items():
            if value is not None:
                if skip_if_set and key in preexisting:
                    if actions is not None:
                        actions.append(("SKIP", key, preexisting[key], str(env_file)))
                else:
                    os.environ[key] = value
                    if actions is not None:
                        actions.append(("SET", key, value, str(env_file)))


def _apply_entries(
    early_config: pytest.Config,
    actions: list[tuple[str, str, str, str]] | None,
) -> None:
    source = _config_source(early_config) if actions is not None else ""
    for entry in _load_values(early_config):
        if entry.unset:
            os.environ.pop(entry.key, None)
            if actions is not None:
                actions.append(("UNSET", entry.key, "", source))
        elif entry.skip_if_set and entry.key in os.environ:
            if actions is not None:
                actions.append(("SKIP", entry.key, os.environ[entry.key], source))
        else:
            final = entry.value.format(**os.environ) if entry.transform else entry.value
            os.environ[entry.key] = final
            if actions is not None:
                actions.append(("SET", entry.key, final, source))


def pytest_report_header(config: pytest.Config) -> list[str] | None:
    """Display environment variable assignments in test session header."""
    if _env_actions_key in config.stash:
        return config.stash[_env_actions_key]
    return None


def _format_actions(actions: list[tuple[str, str, str, str]]) -> list[str]:
    lines = ["pytest-env:"]
    for action, key, value, source in actions:
        if action == "UNSET":
            lines.append(f"  {action:<5} {key}  (from {source})")
        else:
            lines.append(f"  {action:<5} {key}={value}  (from {source})")
    return lines


def _find_toml_config(early_config: pytest.Config) -> Path | None:
    """Find TOML config file by checking inipath first, then walking up the tree."""
    if (
        early_config.inipath
        and early_config.inipath.suffix == ".toml"
        and early_config.inipath.name in {"pytest.toml", ".pytest.toml", "pyproject.toml"}
    ):
        return early_config.inipath

    start_path = early_config.inipath.parent if early_config.inipath is not None else early_config.rootpath
    for current_path in [start_path, *start_path.parents]:
        for toml_name in ("pytest.toml", ".pytest.toml", "pyproject.toml"):
            toml_file = current_path / toml_name
            if toml_file.exists():
                return toml_file
    return None


def _config_source(early_config: pytest.Config) -> str:
    """Describe the configuration source for verbose output."""
    if toml_path := _find_toml_config(early_config):
        _, entries, _ = _load_toml_config(toml_path)
        if entries:
            return str(toml_path)
    if early_config.inipath:
        return str(early_config.inipath)
    return "config"  # pragma: no cover


def _load_toml_config(config_path: Path) -> tuple[list[str], list[Entry], bool | None]:
    """Load env_files and entries from TOML config file."""
    with config_path.open("rb") as file_handler:
        config = tomllib.load(file_handler)

    if config_path.name == "pyproject.toml":
        config = config.get("tool", {})

    pytest_env_config = config.get("pytest_env", {})
    if not pytest_env_config:
        return [], [], None

    raw_env_files = pytest_env_config.get("env_files")
    env_files = [str(f) for f in raw_env_files] if isinstance(raw_env_files, list) else []
    raw_skip = pytest_env_config.get("env_files_skip_if_set")
    env_files_skip_if_set = raw_skip if isinstance(raw_skip, bool) else None

    entries = list(_parse_toml_config(pytest_env_config))
    return env_files, entries, env_files_skip_if_set


def _load_env_files(early_config: pytest.Config, env_files: list[str]) -> Generator[Path, None, None]:
    """Resolve and yield existing env files, with CLI option taking precedence."""
    if cli_envfile := getattr(early_config.known_args_namespace, "envfile", None):
        if cli_envfile.startswith("+"):
            if not (resolved := early_config.rootpath / cli_envfile[1:]).is_file():
                msg = f"Environment file not found: {cli_envfile[1:]}"
                raise FileNotFoundError(msg)
            for env_file_str in env_files or list(early_config.getini("env_files")):
                if (config_resolved := early_config.rootpath / env_file_str).is_file():
                    yield config_resolved
            yield resolved
        else:
            if not (resolved := early_config.rootpath / cli_envfile).is_file():
                msg = f"Environment file not found: {cli_envfile}"
                raise FileNotFoundError(msg)
            yield resolved
        return

    for env_file_str in env_files or list(early_config.getini("env_files")):
        if (resolved := early_config.rootpath / env_file_str).is_file():
            yield resolved


def _load_values(early_config: pytest.Config) -> Iterator[Entry]:
    """Load env entries from config, preferring TOML over INI."""
    if toml_config := _find_toml_config(early_config):
        _, entries, _ = _load_toml_config(toml_config)
        if entries:
            yield from entries
            return

    for line in early_config.getini("env"):
        # INI lines e.g. D:R:NAME=VAL has two flags (R and D), NAME key, and VAL value
        parts = line.partition("=")
        ini_key_parts = parts[0].split(":")
        flags = {k.strip().upper() for k in ini_key_parts[:-1]}
        # R: is a way to designate whether to use raw value -> perform no transformation of the value
        transform = "R" not in flags
        # D: is a way to mark the value to be set only if it does not exist yet
        skip_if_set = "D" in flags
        # U: is a way to unset (remove) an environment variable
        unset = "U" in flags
        key = ini_key_parts[-1].strip()
        value = parts[2].strip()
        yield Entry(key, value, transform, skip_if_set, unset=unset)


def _parse_toml_config(config: dict[str, Any]) -> Generator[Entry, None, None]:
    for key, entry in config.items():
        if key == "env_files" and isinstance(entry, list):
            continue
        if key == "env_files_skip_if_set" and isinstance(entry, bool):
            continue
        if isinstance(entry, dict):
            unset = bool(entry.get("unset"))
            value = str(entry.get("value", "")) if not unset else ""
            transform, skip_if_set = bool(entry.get("transform")), bool(entry.get("skip_if_set"))
        else:
            value, transform, skip_if_set, unset = str(entry), False, False, False
        yield Entry(key, value, transform, skip_if_set, unset=unset)