pristontale-eu-linux/ptdiag/runner.py

"""Structured command logging with argument-array subprocess execution."""

from __future__ import annotations

from dataclasses import dataclass
from datetime import datetime
import json
import os
from pathlib import Path
import subprocess
from typing import Mapping, Sequence


@dataclass(frozen=True)
class CommandResult:
	argv: tuple[str, ...]
	returncode: int | None
	stdout: str
	stderr: str
	timed_out: bool
	dry_run: bool
	error: str | None = None


class CommandLogger:
	def __init__(self, path: Path) -> None:
		self.path = path
		self.path.parent.mkdir(parents=True, exist_ok=True)

	def run(
		self,
		argv: Sequence[str],
		*,
		cwd: Path | None = None,
		env_overrides: Mapping[str, str] | None = None,
		timeout_seconds: float,
		dry_run: bool = False,
	) -> CommandResult:
		arguments = _validate_argv(argv)
		if timeout_seconds <= 0:
			raise ValueError("timeout_seconds must be positive")
		overrides = dict(env_overrides or {})
		started_at = datetime.now().astimezone()
		if dry_run:
			result = CommandResult(arguments, None, "", "", False, True)
			self._append(started_at, cwd, overrides, timeout_seconds, result)
			return result

		environment = os.environ.copy()
		environment.update(overrides)
		try:
			completed = subprocess.run(
				list(arguments),
				cwd=cwd,
				env=environment,
				capture_output=True,
				text=True,
				timeout=timeout_seconds,
				check=False,
			)
			result = CommandResult(
				arguments,
				completed.returncode,
				completed.stdout,
				completed.stderr,
				False,
				False,
			)
		except subprocess.TimeoutExpired as exc:
			result = CommandResult(
				arguments,
				None,
				_decode_timeout_stream(exc.stdout),
				_decode_timeout_stream(exc.stderr),
				True,
				False,
			)
		except OSError as exc:
			result = CommandResult(
				arguments,
				None,
				"",
				"",
				False,
				False,
				str(exc),
			)
		self._append(started_at, cwd, overrides, timeout_seconds, result)
		return result

	def record_lifecycle(
		self,
		argv: Sequence[str],
		*,
		state: str,
		cwd: Path | None,
		env_overrides: Mapping[str, str],
		timeout_seconds: float,
		started_at: datetime,
		pid: int | None = None,
		returncode: int | None = None,
		timed_out: bool = False,
		error: str | None = None,
	) -> None:
		arguments = _validate_argv(argv)
		record = {
			"argv": list(arguments),
			"cwd": str(cwd) if cwd is not None else None,
			"env_overrides": dict(sorted(env_overrides.items())),
			"error": error,
			"observed_at": datetime.now().astimezone().isoformat(),
			"pid": pid,
			"returncode": returncode,
			"started_at": started_at.isoformat(),
			"state": state,
			"timed_out": timed_out,
			"timeout_seconds": timeout_seconds,
		}
		with self.path.open("a", encoding="utf-8") as handle:
			handle.write(json.dumps(record, sort_keys=True))
			handle.write("\n")

	def _append(
		self,
		started_at: datetime,
		cwd: Path | None,
		env_overrides: Mapping[str, str],
		timeout_seconds: float,
		result: CommandResult,
	) -> None:
		record = {
			"argv": list(result.argv),
			"cwd": str(cwd) if cwd is not None else None,
			"dry_run": result.dry_run,
			"env_overrides": dict(sorted(env_overrides.items())),
			"error": result.error,
			"finished_at": datetime.now().astimezone().isoformat(),
			"returncode": result.returncode,
			"started_at": started_at.isoformat(),
			"stderr": result.stderr,
			"stdout": result.stdout,
			"timed_out": result.timed_out,
			"timeout_seconds": timeout_seconds,
		}
		with self.path.open("a", encoding="utf-8") as handle:
			handle.write(json.dumps(record, sort_keys=True))
			handle.write("\n")


def _validate_argv(argv: Sequence[str]) -> tuple[str, ...]:
	if isinstance(argv, (str, bytes)):
		raise TypeError("argv must be a sequence of strings, not a shell command")
	arguments = tuple(argv)
	if not arguments or not all(isinstance(argument, str) and argument for argument in arguments):
		raise ValueError("argv must contain non-empty strings")
	return arguments


def _decode_timeout_stream(value: str | bytes | None) -> str:
	if value is None:
		return ""
	if isinstance(value, bytes):
		return value.decode("utf-8", errors="replace")
	return value