Viewing files at commit 2caecb1e6549Commit detailsView latest

scripts/oakbench/environment.py 101 lines · 3.2 KB · python
1
"""Process environment policy for measured commands.
2
3
All measured subprocesses run through these environments so that author
4
identity, color suppression, and locale-ish noise are controlled in one place.
5
"""
6
7
from __future__ import annotations
8
9
import atexit
10
import os
11
import shutil
12
import tempfile
13
from pathlib import Path, PurePath
14
from typing import Any
15
16
ENV_ISOLATION_VERSION = "1"
17
18
_AUTO_BENCH_HOME: Path | None = None
19
_CLEANUP_REGISTERED = False
20
21
22
def _cleanup_auto_bench_home() -> None:
23
    if _AUTO_BENCH_HOME is not None:
24
        shutil.rmtree(_AUTO_BENCH_HOME, ignore_errors=True)
25
26
27
def _bench_home(home: Path | None = None) -> str:
28
    global _AUTO_BENCH_HOME, _CLEANUP_REGISTERED
29
    if home is None:
30
        path = Path(tempfile.gettempdir()) / f"oak-bench-home-{os.getpid()}"
31
        _AUTO_BENCH_HOME = path
32
        if not _CLEANUP_REGISTERED:
33
            atexit.register(_cleanup_auto_bench_home)
34
            _CLEANUP_REGISTERED = True
35
    else:
36
        path = home
37
    path.mkdir(mode=0o700, parents=True, exist_ok=True)
38
    (path / ".config").mkdir(exist_ok=True)
39
    (path / ".cache").mkdir(exist_ok=True)
40
    return str(path)
41
42
43
def base_env(
44
    author_name: str = "Oak Bench",
45
    author_email: str = "[email protected]",
46
    oak_author: str = "oakbench",
47
    home: Path | None = None,
48
) -> dict[str, str]:
49
    env = os.environ.copy()
50
    bench_home = _bench_home(home)
51
    env.update(
52
        {
53
            "GIT_AUTHOR_NAME": author_name,
54
            "GIT_AUTHOR_EMAIL": author_email,
55
            "GIT_COMMITTER_NAME": author_name,
56
            "GIT_COMMITTER_EMAIL": author_email,
57
            "OAK_AUTHOR": oak_author,
58
            "NO_COLOR": "1",
59
            "CLICOLOR": "0",
60
            "HOME": bench_home,
61
            "XDG_CONFIG_HOME": str(Path(bench_home) / ".config"),
62
            "XDG_CACHE_HOME": str(Path(bench_home) / ".cache"),
63
            "ZDOTDIR": bench_home,
64
            "BASH_ENV": os.devnull,
65
            "ENV": os.devnull,
66
            "GIT_CONFIG_GLOBAL": os.devnull,
67
            "GIT_CONFIG_NOSYSTEM": "1",
68
            "LC_ALL": "C",
69
            "LANG": "C",
70
            "TZ": "UTC",
71
            "OAK_BENCH_ENV_ISOLATION_VERSION": ENV_ISOLATION_VERSION,
72
        }
73
    )
74
    return env
75
76
77
def probe_env(**kwargs: str) -> dict[str, str]:
78
    # The determinism probe intentionally does NOT suppress color: it measures
79
    # the tool's own non-TTY behavior. ANSI bytes in piped output are token
80
    # waste and break provider prompt caching.
81
    env = base_env(**kwargs)
82
    env.pop("NO_COLOR", None)
83
    env.pop("CLICOLOR", None)
84
    return env
85
86
87
def command_display(command: Any) -> str:
88
    """Human/token-countable form of a command or command sequence."""
89
    if (
90
        isinstance(command, list)
91
        and len(command) >= 3
92
        and command[0] in {"/bin/zsh", "zsh", "/bin/bash", "bash", "sh"}
93
        and command[1] in {"-c", "-lc"}
94
    ):
95
        return command[2]
96
    if isinstance(command, list) and command and all(isinstance(item, str) for item in command):
97
        executable = PurePath(command[0]).name if os.sep in command[0] else command[0]
98
        return " ".join([executable, *command[1:]])
99
    if isinstance(command, list):
100
        return " && ".join(command_display(item) for item in command)
101
    return str(command)