cf-tool-cli 1.0.0__py3-none-any.whl

cf_tool/listen.py

@@ -0,0 +1,178 @@
+"""cf listen: problems from the Competitive Companion browser extension, so
+the samples arrive without cf-tool fetching a Codeforces page.
+
+The extension sends each problem as JSON (a POST to localhost) to a fixed
+list of ports (its src/hosts/hosts.ts); cf takes one of them. A contest
+comes as a batch: one POST per problem, each knowing the batch's size. Only
+the extension is listened to: a request from a web page (it carries an
+Origin) or one that isn't JSON is turned away, and so is a page that isn't
+a Codeforces problem.
+"""
+
+from __future__ import annotations
+
+import json
+import socketserver
+from http.server import BaseHTTPRequestHandler, HTTPServer
+from pathlib import Path
+from typing import Callable
+from urllib.parse import urlparse
+
+from cf_tool import ui
+from cf_tool.errors import CFError
+from cf_tool.workspace import (
+    ProblemInfo,
+    Target,
+    load_info,
+    parse_target,
+    problem_folder,
+    read_spec,
+    save_info,
+    save_samples,
+)
+
+DEFAULT_PORT = 4244  # NOTE: Hightail's, a tool gone quiet; the extension still sends there
+REQUEST_TIMEOUT = 5.0  # seconds a connection gets to send its problem
+MAX_BODY = 16 * 1024 * 1024  # bytes
+EXTENSION_ORIGINS = ("chrome-extension://", "moz-extension://", "safari-web-extension://")
+
+# what to do with a problem once its folder is written: (folder, target, info)
+OnProblem = Callable[[Path, Target, ProblemInfo], None]
+
+
+def save_problem(data: dict, root: Path) -> tuple[Path, Target, ProblemInfo, int]:
+    """Write a received problem's tests and info into its folder under root,
+    as cf get would; returns (its folder, the problem, its info, the number
+    of tests). Raises CFError for what isn't a Codeforces problem."""
+    url = str(data.get("url", ""))
+    host = urlparse(url).hostname or ""
+    if host != "codeforces.com" and not host.endswith(".codeforces.com"):
+        raise CFError(f"Not a Codeforces problem: {url or data.get('name', '?')}")
+    parsed = read_spec(url)
+    if not parsed or not parsed.get("problem"):
+        raise CFError(f"Not a Codeforces problem's page: {url}")
+    target = parse_target([url])
+    folder = problem_folder(root, target)
+    tests = [
+        (str(test.get("input", "")), str(test.get("output", "")))
+        for test in data.get("tests", [])
+        if isinstance(test, dict)
+    ]
+    save_samples(folder, tests)
+    limit, memory = data.get("timeLimit"), data.get("memoryLimit")
+    known = load_info(folder)  # NOTE: cf get read more from the statement than the extension sends
+    info = ProblemInfo(
+        name=str(data.get("name", "")),
+        url=url,
+        time_limit=limit / 1000 if isinstance(limit, (int, float)) else None,  # NOTE: sent in ms
+        memory_limit=int(memory) if isinstance(memory, (int, float)) else None,
+        interactive=bool(data.get("interactive", False)),
+    )
+    if known is not None:
+        info.ignore_case, info.epsilon = known.ignore_case, known.epsilon
+        info.any_answer = known.any_answer
+        info.interactive = info.interactive or known.interactive
+        if len(known.cases) == len(tests):
+            info.cases = known.cases
+    save_info(folder, info)
+    return folder, target, info, len(tests)
+
+
+class _Server(HTTPServer):
+    def server_bind(self) -> None:
+        # NOTE: HTTPServer's own also looks up this host's name, which can take
+        # \     half a minute on macOS; the name isn't needed here
+        socketserver.TCPServer.server_bind(self)
+        self.server_name, self.server_port = self.server_address[:2]
+
+
+class _Handler(BaseHTTPRequestHandler):
+    timeout = REQUEST_TIMEOUT  # NOTE: a connection that sends nothing can't hold up the next
+
+    def do_POST(self) -> None:
+        length = self.headers.get("Content-Length") or "0"
+        if not length.isdigit() or int(length) > MAX_BODY:
+            self._answer(413)
+            return
+        body = self.rfile.read(int(length))  # NOTE: read first: an unread body resets the answer
+        origin = self.headers.get("Origin")
+        if origin is not None and not origin.startswith(EXTENSION_ORIGINS):
+            self._answer(403)  # NOTE: a web page (or a sandboxed one: "null"), not the extension
+            return
+        content_type = self.headers.get("Content-Type", "").split(";")[0].strip().lower()
+        if content_type != "application/json":
+            self._answer(415)  # NOTE: a page can't send JSON here without asking first
+            return
+        self._answer(200)
+        self.server.received.append(body)  # type: ignore[attr-defined]
+
+    def _answer(self, status: int) -> None:
+        self.send_response(status)
+        self.end_headers()
+
+    def log_message(self, format, *args) -> None:  # noqa: A002 (the base class's name)
+        pass  # NOTE: no request log on the terminal
+
+
+def listen(
+    root: Path,
+    on_problem: OnProblem | None = None,
+    port: int = DEFAULT_PORT,
+    keep: bool = False,
+    ready: Callable[[int], None] | None = None,
+) -> None:
+    """Receive problems on localhost:port, save each under root and pass it
+    to on_problem (the starter code); returns once a whole batch (a contest,
+    or a single problem) arrived, unless keep. ready(port) is called once
+    listening."""
+    try:
+        server = _Server(("127.0.0.1", port), _Handler)
+    except OSError as e:
+        raise CFError(
+            f"Cannot listen on port {port} ({e.strerror or e}): try another with --port"
+        ) from e
+    server.received = []  # type: ignore[attr-defined]
+    batches: dict[str, set[str]] = {}
+    with server:
+        ui.cyan(
+            f"Listening on port {server.server_port}: click Competitive Companion's button "
+            "on a problem or contest page (Ctrl-C stops)"
+        )
+        if ready is not None:
+            ready(server.server_port)
+        while True:
+            server.handle_request()
+            for body in server.received:  # type: ignore[attr-defined]
+                if _receive(body, root, on_problem, batches) and not keep:
+                    return
+            server.received.clear()  # type: ignore[attr-defined]
+
+
+def _receive(
+    body: bytes, root: Path, on_problem: OnProblem | None, batches: dict[str, set[str]]
+) -> bool:
+    """Save one received problem and say so; returns whether its batch is now complete."""
+    try:
+        data = json.loads(body)
+        if not isinstance(data, dict):
+            raise ValueError("not an object")
+        folder, target, info, count = save_problem(data, root)
+    except (ValueError, TypeError, AttributeError) as e:
+        ui.red(f"Received something that isn't a problem ({e})")
+        return False
+    except CFError as e:
+        ui.red(str(e))
+        return False
+    if on_problem is not None:
+        on_problem(folder, target, info)
+    noun = "test" if count == 1 else "tests"
+    ui.line((f"{target.label}", "bold green"), (f"  {info.name}  {count} {noun} in {folder}", ""))
+    batch = data.get("batch") if isinstance(data.get("batch"), dict) else {}
+    seen = batches.setdefault(str(batch.get("id", "")), set())
+    seen.add(str(data.get("url", "")))
+    size = batch.get("size") if isinstance(batch.get("size"), int) else 1
+    if len(seen) >= size:
+        if size > 1:
+            ui.green(f"All {size} problems are here")
+        return True
+    return False

cf_tool/paths.py

@@ -0,0 +1,57 @@
+"""Where cf-tool keeps its files, and writing them safely.
+
+The folders follow the XDG base directories:
+- settings: $XDG_CONFIG_HOME/cf (config.toml, your starter code in templates/)
+- state: $XDG_STATE_HOME/cf (the login session, caches of what Codeforces said)
+- cache: $XDG_CACHE_HOME/cf (builds, cf-tool's bits/stdc++.h)
+
+Read them as paths.CONFIG_DIR and so on at call time (tests point them elsewhere).
+"""
+
+from __future__ import annotations
+
+import contextlib
+import os
+import tempfile
+from pathlib import Path
+
+
+def _xdg(env_var: str, default: Path) -> Path:
+    value = os.environ.get(env_var, "")
+    return Path(value).expanduser() if value else default
+
+
+CONFIG_DIR = _xdg("XDG_CONFIG_HOME", Path.home() / ".config") / "cf"
+STATE_DIR = _xdg("XDG_STATE_HOME", Path.home() / ".local" / "state") / "cf"
+CACHE_DIR = _xdg("XDG_CACHE_HOME", Path.home() / ".cache") / "cf"
+
+
+def config_file() -> Path:
+    return CONFIG_DIR / "config.toml"
+
+
+def templates_dir() -> Path:
+    """Where your own starter code lives: template.cpp, brute.cpp, gen.py, ..."""
+    return CONFIG_DIR / "templates"
+
+
+def session_file() -> Path:
+    return STATE_DIR / "session"
+
+
+def write_atomic(path: Path, text: str, mode: int = 0o644) -> None:
+    """Replace path with text in one step, created with the given permissions
+    (a Ctrl-C mid-write never leaves a truncated file behind). A symbolic link
+    stays one: the file it points to is written."""
+    path = Path(os.path.realpath(path))  # NOTE: a config.toml kept in a dotfiles repo
+    path.parent.mkdir(parents=True, exist_ok=True)
+    fd, tmp = tempfile.mkstemp(dir=path.parent, prefix=f".{path.name}.")
+    try:
+        with os.fdopen(fd, "w", encoding="utf-8", newline="") as f:
+            f.write(text)
+        os.chmod(tmp, mode)
+        os.replace(tmp, path)
+    except BaseException:
+        with contextlib.suppress(OSError):
+            os.unlink(tmp)
+        raise

cf_tool/process.py

@@ -0,0 +1,431 @@
+"""Running a program: input in, output out, and how it went.
+
+Each program runs in a process group of its own, so a time limit kills
+everything it started (a wrapper script, a forked child) and nothing is left
+running afterwards; Ctrl-C reaches it through cf-tool. Output is read as it
+comes, up to a cap: past it the run stops, the way the judge stops it with
+"Output limit exceeded", instead of filling memory. CPU time and peak memory
+come from the kernel's accounting for the finished program (wait4); Windows
+gives wall time and polled memory.
+"""
+
+from __future__ import annotations
+
+import collections
+import contextlib
+import os
+import shutil
+import signal
+import subprocess
+import sys
+import threading
+import time
+from dataclasses import dataclass
+from pathlib import Path
+from typing import IO, Callable, Sequence
+
+import psutil
+
+from cf_tool.errors import CFError
+
+try:
+    import resource
+except ImportError:  # Windows
+    resource = None  # type: ignore[assignment]
+
+OUTPUT_LIMIT = 64 * 1024 * 1024  # bytes a program may print before it is stopped
+STDERR_KEPT = 64 * 1024  # the last bytes of its error output, kept for messages
+STACK_LIMIT = 256 * 1024 * 1024  # bytes: the stack Codeforces gives
+_CHUNK = 64 * 1024
+_GRACE = 2.0  # seconds to wait for a killed program, and for its pipes to close
+
+OK = "ok"
+EXIT = "exit"  # a non-zero exit code
+SIGNAL = "signal"  # killed by a signal: a crash
+TIMEOUT = "timeout"
+OUTPUT = "output"  # printed more than its output limit
+
+# Windows' codes for the usual crashes, as unsigned exit codes
+_WINDOWS_CRASHES = {
+    0xC0000005: "access violation",
+    0xC00000FD: "stack overflow",
+    0xC0000094: "integer division by zero",
+    0xC0000409: "stack buffer overrun",
+}
+
+
+@dataclass(frozen=True)
+class Result:
+    """How a run went."""
+
+    status: str  # OK, EXIT, SIGNAL, TIMEOUT or OUTPUT
+    code: int | None  # its exit code (negative: the number of the signal that killed it)
+    stdout: bytes
+    stderr: bytes  # its last STDERR_KEPT bytes
+    cpu: float  # seconds of CPU time (wall time where the OS can't tell)
+    wall: float  # seconds
+    memory: int  # peak resident memory in bytes, 0 when unknown
+    timeout: float | None = None  # the seconds it was allowed
+    output_limit: int = OUTPUT_LIMIT
+
+    @property
+    def ok(self) -> bool:
+        return self.status == OK
+
+    def reason(self) -> str:
+        """Why it didn't end well, in a few words: "exit status 1", "killed
+        by SIGSEGV", "killed after 6s", "printed over 64 MB"."""
+        if self.status == TIMEOUT:
+            return f"killed after {self.timeout:g}s"
+        if self.status == OUTPUT:
+            return f"printed over {self.output_limit // (1024 * 1024)} MB"
+        if self.status == SIGNAL and self.code is not None:
+            try:
+                return f"killed by {signal.Signals(-self.code).name}"
+            except ValueError:
+                return f"killed by signal {-self.code}"
+        crash = _WINDOWS_CRASHES.get(self.code or 0)
+        return crash or f"exit status {self.code}"
+
+
+def split_command(command: str) -> list[str]:
+    """The words of a command line. Whitespace separates them; a quote of
+    either kind opens a span that the next quote of either kind closes.
+    Nothing else is special: backslashes stay, so Windows paths survive."""
+    words: list[str] = []
+    word: list[str] = []
+    quoted = False
+    for char in command:
+        if char.isspace() and not quoted:
+            if word:
+                words.append("".join(word))
+                word.clear()
+        elif char in "\"'":
+            if quoted:
+                words.append("".join(word))
+                word.clear()
+            quoted = not quoted
+        else:
+            word.append(char)
+    if word:
+        words.append("".join(word))
+    return words
+
+
+# ┌──────────────────────────────────────────────────────────┐
+# │                        A program                         │
+# └──────────────────────────────────────────────────────────┘
+
+
+class Child:
+    """A started program in a process group of its own, waited for by a
+    thread; `done` is set once it ended, with its exit code, CPU time and
+    peak memory filled in."""
+
+    def __init__(
+        self,
+        argv: Sequence[str],
+        *,
+        cwd: Path,
+        stdin,
+        stdout,
+        stderr,
+        big_stack: bool = False,
+    ) -> None:
+        options: dict = {}
+        if os.name == "posix":
+            options["start_new_session"] = True
+            if big_stack and sys.platform.startswith("linux"):
+                options["preexec_fn"] = _raise_stack
+        else:
+            options["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP
+        argv = list(argv)
+        if os.name == "nt":
+            argv[0] = shutil.which(argv[0]) or argv[0]  # NOTE: finds .bat and .cmd tools too
+        try:
+            self.process = subprocess.Popen(
+                argv, cwd=cwd, stdin=stdin, stdout=stdout, stderr=stderr, **options
+            )
+        except OSError as e:
+            raise CFError(f"Cannot run {argv[0]!r}: {e.strerror or e}") from e
+        self.started = time.perf_counter()
+        self.wall = 0.0
+        self.cpu: float | None = None
+        self.memory = 0
+        self.done = threading.Event()
+        threading.Thread(target=self._wait, daemon=True).start()
+        if os.name != "posix":
+            threading.Thread(target=self._poll_memory, daemon=True).start()
+
+    def _wait(self) -> None:
+        try:
+            if os.name == "posix":
+                _, status, usage = os.wait4(self.process.pid, 0)
+                self.process.returncode = os.waitstatus_to_exitcode(status)
+                self.cpu = usage.ru_utime + usage.ru_stime
+                # NOTE: ru_maxrss is in kilobytes on Linux, in bytes on macOS
+                self.memory = usage.ru_maxrss * (1 if sys.platform == "darwin" else 1024)
+                self._kill_group()  # NOTE: whatever it left running goes with it
+            else:
+                self.process.wait()
+                self.cpu = _windows_cpu(self.process.pid)
+        finally:
+            self.wall = time.perf_counter() - self.started
+            self.done.set()
+
+    def kill(self) -> None:
+        """Kill the program and everything it started."""
+        if os.name == "posix":
+            self._kill_group()
+            return
+        try:
+            parent = psutil.Process(self.process.pid)
+            for child in parent.children(recursive=True):
+                child.kill()
+            parent.kill()
+        except psutil.Error:
+            pass
+
+    def _kill_group(self) -> None:
+        try:
+            os.killpg(self.process.pid, signal.SIGKILL)
+        except (ProcessLookupError, PermissionError):
+            pass  # NOTE: the group is gone already
+
+    def _poll_memory(self) -> None:
+        try:
+            process = psutil.Process(self.process.pid)
+            while not self.done.is_set():
+                info = process.memory_info()
+                self.memory = max(self.memory, getattr(info, "peak_wset", info.rss))
+                time.sleep(0.01)
+        except psutil.Error:
+            pass
+
+
+def _windows_cpu(pid: int) -> float | None:
+    try:
+        times = psutil.Process(pid).cpu_times()
+        return times.user + times.system
+    except psutil.Error:
+        return None
+
+
+def _raise_stack() -> None:
+    """Give the program Codeforces' stack (runs in the child, before exec)."""
+    assert resource is not None
+    soft, hard = resource.getrlimit(resource.RLIMIT_STACK)
+    want = STACK_LIMIT if hard == resource.RLIM_INFINITY else min(STACK_LIMIT, hard)
+    if soft != resource.RLIM_INFINITY and soft < want:
+        resource.setrlimit(resource.RLIMIT_STACK, (want, hard))
+
+
+class _Reader(threading.Thread):
+    """Reads a pipe to its end, keeping its first `limit` bytes (calling
+    over() once there are more) or only its last `tail` bytes."""
+
+    def __init__(
+        self,
+        pipe: IO[bytes],
+        *,
+        limit: int | None = None,
+        tail: int | None = None,
+        over: Callable[[], None] | None = None,
+    ) -> None:
+        super().__init__(daemon=True)
+        self.pipe, self.limit, self.tail, self.over = pipe, limit, tail, over
+        self.chunks: collections.deque[bytes] = collections.deque()
+        self.size = 0
+        self.exceeded = False
+        self.start()
+
+    def run(self) -> None:
+        fd = self.pipe.fileno()
+        try:
+            while chunk := os.read(fd, _CHUNK):
+                if self.exceeded:
+                    continue  # NOTE: still drained, so the program isn't blocked writing
+                self.chunks.append(chunk)
+                self.size += len(chunk)
+                if self.limit is not None and self.size > self.limit:
+                    self.exceeded = True
+                    if self.over is not None:
+                        self.over()
+                while self.tail is not None and self.size - len(self.chunks[0]) >= self.tail:
+                    self.size -= len(self.chunks.popleft())
+        except OSError:
+            pass
+        finally:
+            with contextlib.suppress(OSError):
+                self.pipe.close()
+
+    def data(self) -> bytes:
+        data = b"".join(list(self.chunks))
+        if self.limit is not None:
+            data = data[: self.limit]
+        if self.tail is not None:
+            data = data[-self.tail :]
+        return data
+
+
+def _feed(pipe: IO[bytes], data: bytes) -> None:
+    """Write data to the program's input in the background, then close it."""
+
+    def write() -> None:
+        try:
+            pipe.write(data)
+        except OSError:
+            pass  # NOTE: a program needn't read all of its input
+        finally:
+            with contextlib.suppress(OSError):
+                pipe.close()
+
+    threading.Thread(target=write, daemon=True).start()
+
+
+# ┌──────────────────────────────────────────────────────────┐
+# │                         Running                          │
+# └──────────────────────────────────────────────────────────┘
+
+
+def run(
+    argv: Sequence[str],
+    *,
+    cwd: Path,
+    stdin: bytes | Path | None = None,
+    timeout: float | None,
+    output_limit: int = OUTPUT_LIMIT,
+    big_stack: bool = False,
+) -> Result:
+    """Run argv in cwd and wait for it to end; returns how it went.
+    - stdin: bytes to feed it, a file to read from, or nothing (end of input)
+    - running past timeout seconds, or printing over output_limit bytes,
+      kills it along with everything it started
+    - big_stack: Codeforces' 256 MB stack (where a process limit gives it)"""
+    in_file = stdin.open("rb") if isinstance(stdin, Path) else None
+    feed = stdin if isinstance(stdin, bytes) else None
+    try:
+        child = Child(
+            argv,
+            cwd=cwd,
+            stdin=in_file or (subprocess.PIPE if feed is not None else subprocess.DEVNULL),
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            big_stack=big_stack,
+        )
+    finally:
+        if in_file is not None:
+            in_file.close()  # NOTE: the program has a copy of its own
+    try:
+        out = _Reader(child.process.stdout, limit=output_limit, over=child.kill)
+        err = _Reader(child.process.stderr, tail=STDERR_KEPT)
+        if feed is not None:
+            _feed(child.process.stdin, feed)
+        finished = child.done.wait(timeout)
+        if not finished:
+            child.kill()
+            child.done.wait(_GRACE)
+    except BaseException:  # NOTE: Ctrl-C doesn't reach a program in its own process group
+        child.kill()
+        raise
+    out.join(_GRACE)
+    err.join(_GRACE)
+    return _result(child, finished, out, err, timeout, output_limit)
+
+
+def _result(
+    child: Child,
+    finished: bool,
+    out: _Reader | None,
+    err: _Reader,
+    timeout: float | None,
+    output_limit: int = OUTPUT_LIMIT,
+) -> Result:
+    code = child.process.returncode
+    if out is not None and out.exceeded:
+        status = OUTPUT
+    elif not finished:
+        status = TIMEOUT
+    elif code == 0:
+        status = OK
+    elif code is not None and code < 0:
+        status = SIGNAL
+    else:
+        status = EXIT
+    return Result(
+        status=status,
+        code=code,
+        stdout=out.data() if out is not None else b"",
+        stderr=err.data(),
+        cpu=child.cpu if child.cpu is not None else child.wall,
+        wall=child.wall,
+        memory=child.memory,
+        timeout=timeout,
+        output_limit=output_limit,
+    )
+
+
+def interact(
+    solution: Sequence[str],
+    interactor: Sequence[str],
+    *,
+    cwd: Path,
+    timeout: float,
+    solution_cwd: Path | None = None,
+    big_stack: bool = False,
+) -> tuple[Result, Result]:
+    """Run a solution against an interactor, each one's output the other's
+    input; returns (the solution's result, the interactor's). Steps:
+    - start both, wired together
+    - wait for both to end, at most timeout seconds
+    - once the interactor has ended (its verdict is in), give the solution a
+      moment to end too, then stop it"""
+    to_solution = os.pipe()
+    to_interactor = os.pipe()
+    try:
+        solver = Child(
+            solution,
+            cwd=solution_cwd or cwd,
+            stdin=to_solution[0],
+            stdout=to_interactor[1],
+            stderr=subprocess.PIPE,
+            big_stack=big_stack,
+        )
+        try:
+            judge = Child(
+                interactor,
+                cwd=cwd,
+                stdin=to_interactor[0],
+                stdout=to_solution[1],
+                stderr=subprocess.PIPE,
+            )
+        except BaseException:
+            solver.kill()
+            raise
+    finally:
+        for fd in (*to_solution, *to_interactor):
+            os.close(fd)  # NOTE: only the two programs hold the pipes now
+    solver_err = _Reader(solver.process.stderr, tail=STDERR_KEPT)
+    judge_err = _Reader(judge.process.stderr, tail=STDERR_KEPT)
+    deadline = time.monotonic() + timeout
+    try:
+        while not (solver.done.is_set() and judge.done.is_set()):
+            left = deadline - time.monotonic()
+            if left <= 0:
+                break
+            if judge.done.is_set():
+                solver.done.wait(min(left, _GRACE))  # NOTE: the verdict is in
+                break
+            judge.done.wait(min(left, 0.05))
+    finally:
+        solver_done, judge_done = solver.done.is_set(), judge.done.is_set()
+        for child in (solver, judge):
+            if not child.done.is_set():
+                child.kill()
+                child.done.wait(_GRACE)
+    solver_err.join(_GRACE)
+    judge_err.join(_GRACE)
+    return (
+        _result(solver, solver_done, None, solver_err, timeout),
+        _result(judge, judge_done, None, judge_err, timeout),
+    )