vscode-common-python-lsp 0.1.0__py3-none-any.whl

vscode_common_python_lsp/jsonrpc.py

@@ -0,0 +1,306 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+"""Light-weight JSON-RPC over standard IO."""
+
+from __future__ import annotations
+
+import atexit
+import contextlib
+import json
+import os
+import subprocess
+import threading
+import uuid
+from collections.abc import Sequence
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass
+from typing import BinaryIO
+
+CONTENT_LENGTH = "Content-Length: "
+
+
+class StreamClosedException(Exception):
+    """JSON RPC stream is closed."""
+
+    pass
+
+
+class JsonRpc:
+    """Manages sending and receiving data over JSON-RPC.
+
+    Handles content-length-framed JSON messages over a pair of binary
+    streams (typically ``stdin``/``stdout`` of a subprocess).
+    """
+
+    def __init__(self, reader: BinaryIO, writer: BinaryIO):
+        self._reader = reader
+        self._writer = writer
+        self._write_lock = threading.Lock()
+
+    # -- write --------------------------------------------------------------
+
+    def write(self, data: object) -> None:
+        """Write *data* to the stream in JSON-RPC format (thread-safe)."""
+        with self._write_lock:
+            if self._writer.closed:
+                raise StreamClosedException()
+            content = json.dumps(data)
+            length = len(content.encode("utf-8"))
+            self._writer.write(
+                f"{CONTENT_LENGTH}{length}\r\n\r\n{content}".encode("utf-8")
+            )
+            self._writer.flush()
+
+    def send_data(self, data: object) -> None:
+        """Alias for :meth:`write` — send *data* in JSON-RPC format."""
+        self.write(data)
+
+    # -- read ---------------------------------------------------------------
+
+    def read(self) -> dict[str, object]:
+        """Read and return the next JSON-RPC message from the stream."""
+        if self._reader.closed:
+            raise StreamClosedException()
+        length = None
+        while not length:
+            line = self._readline().decode("utf-8")
+            if line.startswith(CONTENT_LENGTH):
+                length = int(line[len(CONTENT_LENGTH) :])  # noqa: E203
+
+        line = self._readline().decode("utf-8").strip()
+        while line:
+            line = self._readline().decode("utf-8").strip()
+
+        content = self._reader.read(length).decode("utf-8")
+        return json.loads(content)
+
+    def receive_data(self) -> dict[str, object]:
+        """Alias for :meth:`read` — receive data in JSON-RPC format."""
+        return self.read()
+
+    # -- close --------------------------------------------------------------
+
+    def close(self) -> None:
+        """Close both the reader and writer streams."""
+        with contextlib.suppress(Exception):
+            if not self._reader.closed:
+                self._reader.close()
+        with self._write_lock:
+            if not self._writer.closed:
+                self._writer.close()
+
+    # -- internal -----------------------------------------------------------
+
+    def _readline(self) -> bytes:
+        line = self._reader.readline()
+        if not line:
+            raise EOFError
+        return line
+
+
+# ---------------------------------------------------------------------------
+# Process management
+# ---------------------------------------------------------------------------
+
+
+class ProcessManager:
+    """Manages sub-processes launched for running tools."""
+
+    def __init__(self):
+        self._processes: dict[str, subprocess.Popen] = {}
+        self._rpc: dict[str, JsonRpc] = {}
+        self._lock = threading.Lock()
+        self._thread_pool = ThreadPoolExecutor(10)
+
+    def stop_process(self, workspace: str) -> None:
+        """Stop the process for the given workspace."""
+        with self._lock:
+            if workspace in self._processes:
+                proc = self._processes[workspace]
+                try:
+                    proc.kill()
+                    proc.wait(timeout=5)
+                except (OSError, subprocess.TimeoutExpired):
+                    pass
+                del self._processes[workspace]
+            if workspace in self._rpc:
+                rpc = self._rpc.pop(workspace)
+                with contextlib.suppress(Exception):
+                    rpc.close()
+
+    def stop_all_processes(self) -> None:
+        """Stop all managed processes and shutdown transport."""
+        with self._lock:
+            rpcs = list(self._rpc.values())
+            procs = list(self._processes.values())
+            self._rpc.clear()
+            self._processes.clear()
+        for rpc in rpcs:
+            with contextlib.suppress(Exception):
+                rpc.send_data({"id": str(uuid.uuid4()), "method": "exit"})
+            with contextlib.suppress(Exception):
+                rpc.close()
+        for proc in procs:
+            with contextlib.suppress(Exception):
+                proc.kill()
+                proc.wait(timeout=5)
+        self._thread_pool.shutdown(wait=False)
+
+    def start_process(
+        self,
+        workspace: str,
+        args: Sequence[str],
+        cwd: str,
+        env: dict[str, str] | None = None,
+    ) -> None:
+        """Starts a process and establishes JSON-RPC communication over stdio."""
+        new_env = os.environ.copy()
+        if env is not None:
+            new_env.update(env)
+        proc = subprocess.Popen(
+            args,
+            cwd=cwd,
+            stdout=subprocess.PIPE,
+            stdin=subprocess.PIPE,
+            env=new_env,
+        )
+        with self._lock:
+            self._processes[workspace] = proc
+            self._rpc[workspace] = JsonRpc(proc.stdout, proc.stdin)
+
+        def _monitor_process():
+            proc.wait()
+            with self._lock:
+                # Only clean up if this is still the registered process
+                # (a replacement may have been started under the same key).
+                if self._processes.get(workspace) is not proc:
+                    return
+                self._processes.pop(workspace, None)
+                rpc = self._rpc.pop(workspace, None)
+                if rpc:
+                    rpc.close()
+
+        self._thread_pool.submit(_monitor_process)
+
+    def get_json_rpc(self, workspace: str) -> JsonRpc:
+        """Gets the JSON-RPC wrapper for a given workspace id."""
+        with self._lock:
+            if workspace in self._rpc:
+                return self._rpc[workspace]
+        raise StreamClosedException()
+
+
+_process_manager = ProcessManager()
+_start_lock = threading.Lock()
+atexit.register(_process_manager.stop_all_processes)
+
+
+def _get_json_rpc(workspace: str) -> JsonRpc | None:
+    try:
+        return _process_manager.get_json_rpc(workspace)
+    except (StreamClosedException, KeyError):
+        return None
+
+
+def get_or_start_json_rpc(
+    workspace: str,
+    interpreter: Sequence[str],
+    cwd: str,
+    runner_script: str,
+    env: dict[str, str] | None = None,
+) -> JsonRpc | None:
+    """Gets an existing JSON-RPC connection or starts one and return it."""
+    with _start_lock:
+        res = _get_json_rpc(workspace)
+        if not res:
+            args = [*interpreter, runner_script]
+            _process_manager.start_process(workspace, args, cwd, env)
+            res = _get_json_rpc(workspace)
+    return res
+
+
+# ---------------------------------------------------------------------------
+# RPC run helpers
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class RpcRunResult:
+    """Object to hold result from running tool over RPC."""
+
+    stdout: str
+    stderr: str
+    exception: str | None = None
+
+
+def run_over_json_rpc(
+    workspace: str,
+    interpreter: Sequence[str],
+    module: str,
+    argv: Sequence[str],
+    use_stdin: bool,
+    cwd: str,
+    runner_script: str,
+    source: str | None = None,
+    env: dict[str, str] | None = None,
+    timeout: float | None = None,
+) -> RpcRunResult:
+    """Uses JSON-RPC to execute a command."""
+    rpc = get_or_start_json_rpc(workspace, interpreter, cwd, runner_script, env)
+    if not rpc:
+        raise ConnectionError("Failed to run over JSON-RPC.")
+
+    msg_id = str(uuid.uuid4())
+    msg = {
+        "id": msg_id,
+        "method": "run",
+        "module": module,
+        "argv": argv,
+        "useStdin": use_stdin,
+        "cwd": cwd,
+    }
+    if source is not None:
+        msg["source"] = source
+
+    rpc.send_data(msg)
+
+    if timeout is not None:
+        result_container: list[object] = [None]
+        error_container: list[BaseException | None] = [None]
+
+        def _receive():
+            try:
+                result_container[0] = rpc.receive_data()
+            except Exception as e:
+                error_container[0] = e
+
+        recv_thread = threading.Thread(target=_receive, daemon=True)
+        recv_thread.start()
+        recv_thread.join(timeout)
+        if recv_thread.is_alive():
+            _process_manager.stop_process(workspace)
+            raise TimeoutError(f"JSON-RPC call timed out after {timeout}s")
+        if error_container[0] is not None:
+            raise error_container[0]
+        data = result_container[0]
+    else:
+        data = rpc.receive_data()
+
+    if data["id"] != msg_id:
+        return RpcRunResult(
+            "", f"Invalid result for request: {json.dumps(msg, indent=4)}"
+        )
+
+    result = data.get("result", "")
+    if "error" in data:
+        error = data["error"]
+        if data.get("exception", False):
+            return RpcRunResult(result, "", error)
+        return RpcRunResult(result, error)
+
+    return RpcRunResult(result, "")
+
+
+def shutdown_json_rpc():
+    """Shutdown all JSON-RPC processes."""
+    _process_manager.stop_all_processes()

vscode_common_python_lsp/linting.py

@@ -0,0 +1,62 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+"""Lint request tracking for stale-result deduplication.
+
+When multiple ``didSave`` events fire in quick succession each one spawns
+a tool process.  Only the **latest** result should be published — earlier
+runs are stale.  This module provides :class:`LintRequestTracker` which
+manages a per-URI version counter protected by a threading lock.
+
+Shared by flake8, mypy, and pylint (formatters don't need this).
+"""
+
+from __future__ import annotations
+
+import threading
+
+
+class LintRequestTracker:
+    """Thread-safe version counter for lint request deduplication.
+
+    Usage::
+
+        tracker = LintRequestTracker()
+
+        # At the start of a lint run:
+        version = tracker.increment(uri)
+
+        # ... run tool, parse output ...
+
+        # Before publishing diagnostics:
+        if tracker.is_current(uri, version):
+            publish_diagnostics(...)
+    """
+
+    def __init__(self) -> None:
+        self._versions: dict[str, int] = {}
+        self._lock = threading.Lock()
+
+    def increment(self, uri: str) -> int:
+        """Bump and return the version for *uri*."""
+        with self._lock:
+            version = self._versions.get(uri, 0) + 1
+            self._versions[uri] = version
+            return version
+
+    def is_current(self, uri: str, version: int) -> bool:
+        """Return *True* if *version* is still the latest for *uri*."""
+        with self._lock:
+            if uri not in self._versions:
+                return False
+            return self._versions[uri] == version
+
+    def reset(self, uri: str | None = None) -> None:
+        """Reset version counters.
+
+        When *uri* is ``None`` all counters are cleared.
+        """
+        with self._lock:
+            if uri is None:
+                self._versions.clear()
+            else:
+                self._versions.pop(uri, None)

vscode_common_python_lsp/notebook.py

@@ -0,0 +1,245 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+"""Notebook-specific helpers for whole-notebook linting with cross-cell context.
+
+Provides:
+* :class:`SyntheticDocument` — a typed stand-in for ``TextDocument`` used
+  when linting a combined notebook source.
+* :func:`build_notebook_source` — concatenates all code cells into one
+  string, replacing IPython magic lines with ``pass``.
+* :func:`remap_diagnostics_to_cells` — maps combined-source diagnostics
+  back to individual cell URIs.
+* :data:`NOTEBOOK_SYNC_OPTIONS` — default ``NotebookDocumentSyncOptions``
+  shared by all tool extensions.
+
+This module is identical across flake8, isort, and pylint (black uses a
+simpler variant).
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import re
+from collections.abc import Callable, Sequence
+from typing import Any, Protocol
+
+import lsprotocol.types as lsp
+
+# ---------------------------------------------------------------------------
+# Protocols & types
+# ---------------------------------------------------------------------------
+
+
+class TextDocumentLike(Protocol):
+    """Protocol for objects that provide text document attributes.
+
+    Any object with ``source`` and ``language_id`` string attributes
+    satisfies this — including ``pygls.workspace.TextDocument``.
+    """
+
+    source: str
+    language_id: str
+
+
+class CellLike(Protocol):
+    """Protocol for notebook cell objects.
+
+    Any object with ``kind`` and ``document`` attributes satisfies this —
+    including ``lsprotocol.types.NotebookCell``.
+    """
+
+    kind: Any
+    document: str | None
+
+
+@dataclasses.dataclass
+class SyntheticDocument:
+    """Typed stand-in for ``workspace.TextDocument`` used in notebook linting.
+
+    Replaces ``types.SimpleNamespace`` so that the synthetic document has
+    an explicit, portable shape that can be type-checked.
+    """
+
+    uri: str
+    path: str
+    source: str
+    language_id: str = "python"
+    version: int = 0
+
+
+# Matches IPython magic lines (%, %%, !, !!) so they can be replaced with ``pass``.
+MAGIC_LINE_RE = re.compile(r"^\s*(?:%%\w|%(?!=)\w|!!|!(?!=)\w)")
+
+# Maximum character value in LSP (the LSP ``uinteger`` max per the spec).
+# https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#uinteger
+MAX_LSP_CHARACTER = 2_147_483_647
+
+NOTEBOOK_SYNC_OPTIONS = lsp.NotebookDocumentSyncOptions(
+    notebook_selector=[
+        lsp.NotebookDocumentFilterWithNotebook(
+            notebook="jupyter-notebook",
+            cells=[
+                lsp.NotebookCellLanguage(language="python"),
+            ],
+        ),
+        lsp.NotebookDocumentFilterWithNotebook(
+            notebook="interactive",
+            cells=[
+                lsp.NotebookCellLanguage(language="python"),
+            ],
+        ),
+    ],
+    save=True,
+)
+
+
+@dataclasses.dataclass
+class CellOffset:
+    """Describes where a single notebook cell's lines begin in the combined source."""
+
+    cell_uri: str
+    start_line: int
+    line_count: int
+
+
+def build_notebook_source(
+    cells: Sequence[CellLike],
+    get_text_document: Callable[[str], TextDocumentLike | None],
+    *,
+    sanitize_line: Callable[[str], str] | None = None,
+) -> tuple[str, list[CellOffset]]:
+    """Build a single Python source string from all code cells.
+
+    Parameters
+    ----------
+    cells:
+        The notebook's cell list (``nb.cells``).
+    get_text_document:
+        A callable that resolves a cell document URI to a text document
+        object (with ``.source`` and ``.language_id`` attributes).
+    sanitize_line:
+        Optional per-line transformation.  When *None* the default
+        behaviour replaces IPython magic lines with ``pass\\n``.  Pass a
+        custom callable to override (e.g. for tool-specific magic handling).
+
+    Returns
+    -------
+    (combined_source, cell_map) where *cell_map* is a list of
+    :class:`CellOffset` instances describing where each cell's lines
+    begin in the combined source.
+    """
+    _sanitize = sanitize_line or _default_sanitize_line
+
+    source_parts: list[str] = []
+    cell_map: list[CellOffset] = []
+    current_line = 0
+
+    for cell in cells:
+        if cell.kind != lsp.NotebookCellKind.Code or cell.document is None:
+            continue
+        doc = get_text_document(cell.document)
+        if doc is None or doc.language_id != "python":
+            continue
+
+        source = doc.source
+        if not source:
+            continue
+
+        lines = source.splitlines(keepends=True)
+        # Ensure the last line ends with a newline.
+        if lines and not lines[-1].endswith("\n"):
+            lines[-1] += "\n"
+
+        sanitized_lines = [_sanitize(line) for line in lines]
+
+        cell_map.append(CellOffset(cell.document, current_line, len(sanitized_lines)))
+        source_parts.extend(sanitized_lines)
+        current_line += len(sanitized_lines)
+
+    return "".join(source_parts), cell_map
+
+
+def _default_sanitize_line(line: str) -> str:
+    """Replace IPython magic lines with ``pass`` while preserving line endings."""
+    if not MAGIC_LINE_RE.match(line):
+        return line
+    if line.endswith("\r\n"):
+        return "pass\r\n"
+    if line.endswith("\n"):
+        return "pass\n"
+    return "pass"
+
+
+def get_cell_for_line(
+    global_line: int, cell_map: list[CellOffset]
+) -> CellOffset | None:
+    """Return the :class:`CellOffset` entry that owns *global_line*.
+
+    *global_line* is a 0-based line number in the combined notebook source.
+    Returns ``None`` if no cell owns the line.
+    """
+    for entry in cell_map:
+        if entry.start_line <= global_line < entry.start_line + entry.line_count:
+            return entry
+    return None
+
+
+def remap_diagnostics_to_cells(
+    diagnostics: Sequence[lsp.Diagnostic],
+    cell_map: list[CellOffset],
+) -> dict[str, list[lsp.Diagnostic]]:
+    """Map combined-source diagnostics back to individual cell URIs.
+
+    Each diagnostic's line range is adjusted relative to the owning cell.
+    Diagnostics whose start line doesn't fall in any cell are discarded.
+    If a diagnostic's end line crosses a cell boundary it is clamped.
+    """
+    per_cell: dict[str, list[lsp.Diagnostic]] = {
+        entry.cell_uri: [] for entry in cell_map
+    }
+
+    for diag in diagnostics:
+        entry = get_cell_for_line(diag.range.start.line, cell_map)
+        if entry is None:
+            continue
+
+        local_start_line = diag.range.start.line - entry.start_line
+        local_start = lsp.Position(
+            line=local_start_line,
+            character=diag.range.start.character,
+        )
+
+        # Clamp end line to the cell boundary (defensive).
+        max_end_line = entry.line_count - 1
+        raw_end_line = diag.range.end.line - entry.start_line
+        clamped = raw_end_line > max_end_line
+        local_end_line = min(raw_end_line, max_end_line)
+        # Use max LSP uint32 value to cover the full last line.
+        local_end = lsp.Position(
+            line=local_end_line,
+            character=MAX_LSP_CHARACTER if clamped else diag.range.end.character,
+        )
+
+        # Ensure end is not before start (inverted range violates LSP spec).
+        if (
+            local_end.line == local_start.line
+            and local_end.character < local_start.character
+        ):
+            local_end = lsp.Position(
+                line=local_start.line, character=local_start.character
+            )
+
+        remapped = lsp.Diagnostic(
+            range=lsp.Range(start=local_start, end=local_end),
+            message=diag.message,
+            severity=diag.severity,
+            code=diag.code,
+            code_description=diag.code_description,
+            source=diag.source,
+            related_information=diag.related_information,
+            tags=diag.tags,
+            data=diag.data,
+        )
+        per_cell[entry.cell_uri].append(remapped)
+
+    return per_cell