Class WorkerClient

mypy/build.py:267–323

A simple class that represents a mypy build worker.

class WorkerClient:
    """A simple class that represents a mypy build worker."""

    conn: IPCClient

    def __init__(self, status_file: str, options_data: str, env: Mapping[str, str]) -> None:
        self.status_file = status_file
        if os.path.isfile(status_file):
            os.unlink(status_file)

        command = [
            sys.executable,
            "-m",
            "mypy.build_worker",
            f"--status-file={status_file}",
            f"--options-data={options_data}",
        ]
        # Return early without waiting, caller must call connect() before using the client.
        self.proc = subprocess.Popen(command, env=env)
        self.connected = False

    def connect(self) -> None:
        end_time = time.time() + WORKER_START_TIMEOUT
        last_exception: Exception | None = None
        while time.time() < end_time:
            try:
                data = read_status(self.status_file)
            except BadStatus as exc:
                last_exception = exc
                time.sleep(WORKER_START_INTERVAL)
                continue
            try:
                pid, connection_name = data["pid"], data["connection_name"]
                assert isinstance(pid, int), f"Bad PID: {pid}"
                assert isinstance(connection_name, str), f"Bad connection name: {connection_name}"
                if sys.platform != "win32":
                    # Windows uses "wrapper processes" to run Python, so we cannot
                    # verify PIDs reliably.
                    assert pid == self.proc.pid, f"PID mismatch: {pid} vs {self.proc.pid}"
                self.conn = IPCClient(connection_name, WORKER_CONNECTION_TIMEOUT)
                self.connected = True
                return
            except Exception as exc:
                last_exception = exc
                break
        print(f"Failed to establish connection with worker: {last_exception}")

    def close(self) -> None:
        if self.connected:
            self.conn.close()
        # Technically we don't need to wait, but otherwise we will get ResourceWarnings.
        try:
            self.proc.wait(timeout=WORKER_SHUTDOWN_TIMEOUT)
        except subprocess.TimeoutExpired:
            pass
        if os.path.isfile(self.status_file):
            os.unlink(self.status_file)