Registry¶

@classmethod
def from_path(cls, path: Path) -> "OutputKind":
    """The kind a file's suffix implies, text when unknown."""
    match path.suffix.lower():
        case ".csv" | ".parquet":
            return cls.TABLE
        case ".html" | ".htm":
            return cls.HTML
        case ".png" | ".jpg" | ".jpeg" | ".svg" | ".gif" | ".webp":
            return cls.IMAGE
        case ".md":
            return cls.MARKDOWN
        case _:
            return cls.TEXT

@classmethod
def at(cls, path: Path, title: str | None = None) -> "OutputSpec":
    """An output typed from its suffix and titled by its name unless given."""
    return cls(path=path, title=title or path.name, kind=OutputKind.from_path(path))

def modified(self) -> datetime | None:
    """When the file was last written, None while it does not exist."""
    if not self.path.exists():
        return None

    return datetime.fromtimestamp(self.path.stat().st_mtime)

def last_run(self) -> datetime | None:
    """The newest output timestamp, None when nothing has been produced."""
    stamps = [stamp for stamp in (output.modified() for output in self.outputs) if stamp is not None]
    if not stamps:
        return None

    return max(stamps)

def command(self) -> list[str]:
    """The subprocess that runs this spec: the module for a script, the workbench CLI for the rest."""
    if self.kind is WorkbenchKind.SCRIPT:
        return [sys.executable, "-m", self.module, *self.args]

    return [sys.executable, "-m", CLI_MODULE, self.name, *self.args]

def load(self) -> object:
    """Imports and returns what the entry point names."""
    ...

def row(self) -> dict[str, object]:
    """The task as a row of the client state."""
    return {
        TaskField.NAME: self.name,
        TaskField.PID: self.pid,
        TaskField.STARTED: self.started.isoformat(timespec="seconds"),
        TaskField.STATUS: str(self.status),
        TaskField.RETURNCODE: self.returncode,
    }

def running(self, name: str) -> bool:
    """Whether the named task runs."""
    task = self.tasks.get(name)

    return task is not None and task.status is TaskStatus.RUNNING

def rows(self) -> list[dict[str, object]]:
    """Every task as a row of the client state."""
    return [task.row() for task in self.tasks.values()]

async def run(self, name: str, command: Sequence[str], *, cwd: Path | None = None) -> Task:
    """Runs the command to completion, streaming its lines; a name still running is an error.

    Args:
        name: The task name, unique among the running tasks.
        command: The command and its arguments.
        cwd: The directory the command runs in; None keeps the server's.
    """
    if self.running(name):
        raise RuntimeError(f"{name} is already running")

    process = await asyncio.create_subprocess_exec(
        *command,
        cwd=cwd,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.STDOUT,
        env={**os.environ, **STREAM_ENV},
        limit=LINE_LIMIT,
        creationflags=CREATION_FLAGS,
        start_new_session=NEW_SESSION,
    )
    tree = ProcessTree.adopt(process.pid)
    task = Task(name, process, tree)
    self.tasks[name] = task
    log.bind(task=name, pid=task.pid, command=list(command)).info("Script started.")
    self.on_change(task)

    stdout = process.stdout
    if stdout is None:
        raise RuntimeError(f"{name} has no output pipe")

    try:
        async for raw in stdout:
            self.on_line(name, raw.decode("utf-8", errors="replace").rstrip("\r\n"))

        task.returncode = await process.wait()
    finally:
        tree.close()

    if task.stopping:
        task.status = TaskStatus.STOPPED
    else:
        task.status = TaskStatus.DONE if task.returncode == 0 else TaskStatus.FAILED

    log.bind(task=name, returncode=task.returncode, status=str(task.status)).info("Script exited.")
    self.on_change(task)

    return task

async def stop(self, name: str) -> None:
    """Kills the task's process tree; its run then ends as stopped."""
    task = self.tasks.get(name)
    if task is None or task.status is not TaskStatus.RUNNING:
        return

    task.stopping = True
    signalled = await asyncio.to_thread(task.tree.kill)
    log.bind(task=name, processes=signalled).info("Script stopped.")

@classmethod
def adopt(cls, pid: int) -> "ProcessTree":
    """Takes a process spawned with `CREATION_FLAGS` and `NEW_SESSION`: into a new job and let run on Windows, checked to lead its own group elsewhere."""
    if not WINDOWS:
        if os.getpgid(pid) != pid:
            raise ValueError(f"process {pid} does not lead its own group")

        return cls(pid)

    process = _checked("OpenProcess", _SET_QUOTA_AND_TERMINATE, False, pid)
    try:
        job = _checked("CreateJobObjectW", None, None)
        limits = _ExtendedLimits()
        limits.BasicLimitInformation.LimitFlags = _KILL_ON_JOB_CLOSE
        _ = _checked("SetInformationJobObject", job, _EXTENDED_LIMITS, ctypes.byref(limits), ctypes.sizeof(limits))
        _ = _checked("AssignProcessToJobObject", job, process)
    except OSError:
        _ = _win32("TerminateProcess", process, 1)
        raise
    finally:
        _ = _win32("CloseHandle", process)

    psutil.Process(pid).resume()

    return cls(pid, job)

def kill(self, grace: float = KILL_GRACE) -> int:
    """Ends every process in the tree and returns how many there were; POSIX asks first and kills what is left after `grace` seconds, Windows has no gentle signal."""
    if WINDOWS:
        accounting = _BasicAccounting()
        _ = _checked("QueryInformationJobObject", self.job, _BASIC_ACCOUNTING, ctypes.byref(accounting), ctypes.sizeof(accounting), None)
        _ = _checked("TerminateJobObject", self.job, 1)

        return accounting.ActiveProcesses

    held = len(self._members())
    if held == 0:
        return 0

    with suppress(ProcessLookupError):
        os.killpg(self.pid, signal.SIGTERM)

    deadline = time.monotonic() + grace
    while self._members() and time.monotonic() < deadline:
        time.sleep(POLL)

    with suppress(ProcessLookupError):
        os.killpg(self.pid, signal.SIGKILL)

    return held

def close(self) -> None:
    """Ends anything still in the tree and releases it: a job dies with its handle, a group gets SIGKILL."""
    if WINDOWS:
        _ = _win32("CloseHandle", self.job)

        return

    if self._members():
        with suppress(ProcessLookupError):
            os.killpg(self.pid, signal.SIGKILL)