BIG BIG BIG BIG BIG UPDATES
Co-authored-by: Copilot <[email protected]>
This commit is contained in:
gabdevele
co-authored by
Copilot
Unverified
Copilot
parent
46f7cd6a5f
commit
93ba8e9ac2
@@ -0,0 +1,323 @@
|
||||
import asyncio
|
||||
import json
|
||||
import shutil
|
||||
import tempfile
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
WORKSPACE_ROOT = Path(tempfile.gettempdir()) / "theplug-lsp"
|
||||
CONFIG_FILE = "pyrightconfig.json"
|
||||
ENTRY_FILE = "proxy_addon.py"
|
||||
|
||||
STUB_FILES = {
|
||||
"typings/mitmproxy/__init__.pyi": """from . import (
|
||||
command,
|
||||
ctx,
|
||||
dns,
|
||||
http,
|
||||
tcp,
|
||||
tls,
|
||||
udp,
|
||||
websocket,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"command",
|
||||
"ctx",
|
||||
"dns",
|
||||
"http",
|
||||
"tcp",
|
||||
"tls",
|
||||
"udp",
|
||||
"websocket",
|
||||
]
|
||||
""",
|
||||
"typings/mitmproxy/ctx.pyi": """class _Logger:
|
||||
def debug(self, message: object) -> None: ...
|
||||
def info(self, message: object) -> None: ...
|
||||
def warn(self, message: object) -> None: ...
|
||||
def error(self, message: object) -> None: ...
|
||||
def alert(self, message: object) -> None: ...
|
||||
|
||||
|
||||
class _Options:
|
||||
def __getattr__(self, name: str) -> object: ...
|
||||
|
||||
|
||||
log: _Logger
|
||||
options: _Options
|
||||
""",
|
||||
"typings/mitmproxy/http.pyi": """from typing import Any
|
||||
|
||||
|
||||
class Headers(dict[str, str]):
|
||||
...
|
||||
|
||||
|
||||
class Request:
|
||||
host: str
|
||||
path: str
|
||||
headers: Headers
|
||||
query: dict[str, str]
|
||||
content: bytes
|
||||
|
||||
|
||||
class Response:
|
||||
headers: Headers
|
||||
status_code: int
|
||||
content: bytes
|
||||
|
||||
@classmethod
|
||||
def make(
|
||||
cls,
|
||||
status_code: int,
|
||||
content: bytes | str = b"",
|
||||
headers: dict[str, str] | None = None,
|
||||
) -> "Response": ...
|
||||
|
||||
|
||||
class WebSocketMessage:
|
||||
content: bytes
|
||||
from_client: bool
|
||||
|
||||
|
||||
class WebSocketData:
|
||||
messages: list[WebSocketMessage]
|
||||
|
||||
|
||||
class HTTPFlow:
|
||||
request: Request
|
||||
response: Response | None
|
||||
websocket: WebSocketData | None
|
||||
metadata: dict[str, Any]
|
||||
""",
|
||||
"typings/mitmproxy/tcp.pyi": """class TCPMessage:
|
||||
content: bytes
|
||||
from_client: bool
|
||||
|
||||
|
||||
class TCPFlow:
|
||||
messages: list[TCPMessage]
|
||||
""",
|
||||
"typings/mitmproxy/dns.pyi": """class Question:
|
||||
name: str
|
||||
type: int
|
||||
|
||||
|
||||
class ResourceRecord:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
name: str,
|
||||
type: int,
|
||||
class_: int,
|
||||
ttl: int,
|
||||
data: bytes,
|
||||
) -> None: ...
|
||||
|
||||
|
||||
class Message:
|
||||
question: Question | None
|
||||
|
||||
def succeed(self, records: list[ResourceRecord]) -> "Message": ...
|
||||
|
||||
|
||||
class _Types:
|
||||
A: int
|
||||
|
||||
|
||||
class _Classes:
|
||||
IN: int
|
||||
|
||||
|
||||
types: _Types
|
||||
classes: _Classes
|
||||
|
||||
|
||||
class DNSFlow:
|
||||
request: Message
|
||||
response: Message | None
|
||||
""",
|
||||
"typings/mitmproxy/command.pyi": """from collections.abc import Callable
|
||||
from typing import TypeVar
|
||||
|
||||
F = TypeVar("F", bound=Callable[..., object])
|
||||
|
||||
|
||||
def command(name: str) -> Callable[[F], F]: ...
|
||||
""",
|
||||
"typings/mitmproxy/websocket.pyi": """class WebSocketMessage:
|
||||
content: bytes
|
||||
from_client: bool
|
||||
""",
|
||||
"typings/mitmproxy/tls.pyi": """class ClientHelloData: ...
|
||||
""",
|
||||
"typings/mitmproxy/udp.pyi": """class UDPMessage:
|
||||
content: bytes
|
||||
from_client: bool
|
||||
""",
|
||||
}
|
||||
|
||||
|
||||
class LspFrameBuffer:
|
||||
def __init__(self) -> None:
|
||||
self.buffer = bytearray()
|
||||
|
||||
def feed(self, chunk: bytes) -> list[str]:
|
||||
self.buffer.extend(chunk)
|
||||
messages: list[str] = []
|
||||
while True:
|
||||
header_end = self.buffer.find(b"\r\n\r\n")
|
||||
if header_end < 0:
|
||||
return messages
|
||||
|
||||
header = self.buffer[:header_end].decode("ascii", errors="ignore")
|
||||
size = content_length(header)
|
||||
if size is None:
|
||||
del self.buffer[: header_end + 4]
|
||||
continue
|
||||
|
||||
frame_end = header_end + 4 + size
|
||||
if len(self.buffer) < frame_end:
|
||||
return messages
|
||||
|
||||
body = bytes(self.buffer[header_end + 4 : frame_end])
|
||||
del self.buffer[:frame_end]
|
||||
messages.append(body.decode("utf-8"))
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class PythonEditorSession:
|
||||
root_path: Path
|
||||
file_path: Path
|
||||
|
||||
@property
|
||||
def root_uri(self) -> str:
|
||||
return self.root_path.as_uri()
|
||||
|
||||
@property
|
||||
def file_uri(self) -> str:
|
||||
return self.file_path.as_uri()
|
||||
|
||||
|
||||
def content_length(header: str) -> int | None:
|
||||
for line in header.split("\r\n"):
|
||||
if ":" not in line:
|
||||
continue
|
||||
name, value = line.split(":", 1)
|
||||
if name.lower() == "content-length":
|
||||
return int(value.strip())
|
||||
return None
|
||||
|
||||
|
||||
def encode_message(message: str) -> bytes:
|
||||
payload = message.encode("utf-8")
|
||||
header = f"Content-Length: {len(payload)}\r\n\r\n".encode("ascii")
|
||||
return header + payload
|
||||
|
||||
|
||||
def lsp_binary() -> str | None:
|
||||
return shutil.which("basedpyright-langserver")
|
||||
|
||||
|
||||
def get_session() -> PythonEditorSession:
|
||||
root_path = WORKSPACE_ROOT
|
||||
file_path = root_path / ENTRY_FILE
|
||||
|
||||
root_path.mkdir(parents=True, exist_ok=True)
|
||||
file_path.touch(exist_ok=True)
|
||||
|
||||
config = {
|
||||
"include": [ENTRY_FILE],
|
||||
"stubPath": "typings",
|
||||
"typeCheckingMode": "basic",
|
||||
"pythonVersion": "3.14",
|
||||
"reportMissingTypeStubs": False,
|
||||
"reportMissingImports": "warning",
|
||||
"useLibraryCodeForTypes": True,
|
||||
}
|
||||
(root_path / CONFIG_FILE).write_text(
|
||||
json.dumps(config, indent=2),
|
||||
encoding="utf-8",
|
||||
)
|
||||
|
||||
for relative_path, content in STUB_FILES.items():
|
||||
path = root_path / relative_path
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
path.write_text(content, encoding="utf-8")
|
||||
|
||||
return PythonEditorSession(root_path=root_path, file_path=file_path)
|
||||
|
||||
|
||||
async def open_process(session: PythonEditorSession) -> asyncio.subprocess.Process:
|
||||
binary = lsp_binary()
|
||||
if binary is None:
|
||||
raise RuntimeError("basedpyright-langserver is not installed")
|
||||
|
||||
return await asyncio.create_subprocess_exec(
|
||||
binary,
|
||||
"--stdio",
|
||||
cwd=session.root_path,
|
||||
stdin=asyncio.subprocess.PIPE,
|
||||
stdout=asyncio.subprocess.PIPE,
|
||||
stderr=asyncio.subprocess.PIPE,
|
||||
)
|
||||
|
||||
|
||||
async def forward_client(websocket, writer: asyncio.StreamWriter) -> None:
|
||||
while True:
|
||||
message = await websocket.receive_text()
|
||||
writer.write(encode_message(message))
|
||||
await writer.drain()
|
||||
|
||||
|
||||
async def forward_server(reader: asyncio.StreamReader, websocket) -> None:
|
||||
frames = LspFrameBuffer()
|
||||
while True:
|
||||
chunk = await reader.read(65536)
|
||||
if not chunk:
|
||||
return
|
||||
for message in frames.feed(chunk):
|
||||
await websocket.send_text(message)
|
||||
|
||||
|
||||
async def drain_errors(reader: asyncio.StreamReader) -> None:
|
||||
while True:
|
||||
chunk = await reader.read(4096)
|
||||
if not chunk:
|
||||
return
|
||||
|
||||
|
||||
async def bridge(websocket, session: PythonEditorSession) -> None:
|
||||
process = await open_process(session)
|
||||
stdin = process.stdin
|
||||
stdout = process.stdout
|
||||
stderr = process.stderr
|
||||
|
||||
if stdin is None or stdout is None or stderr is None:
|
||||
process.kill()
|
||||
await process.wait()
|
||||
raise RuntimeError("basedpyright-langserver did not expose stdio pipes")
|
||||
|
||||
send_task = asyncio.create_task(forward_client(websocket, stdin))
|
||||
receive_task = asyncio.create_task(forward_server(stdout, websocket))
|
||||
error_task = asyncio.create_task(drain_errors(stderr))
|
||||
|
||||
try:
|
||||
done, pending = await asyncio.wait(
|
||||
{send_task, receive_task, error_task},
|
||||
return_when=asyncio.FIRST_COMPLETED,
|
||||
)
|
||||
for task in pending:
|
||||
task.cancel()
|
||||
for task in done:
|
||||
task.result()
|
||||
finally:
|
||||
if process.returncode is None:
|
||||
stdin.close()
|
||||
await stdin.wait_closed()
|
||||
try:
|
||||
await asyncio.wait_for(process.wait(), timeout=1)
|
||||
except TimeoutError:
|
||||
process.kill()
|
||||
await process.wait()
|
||||
Reference in New Issue
Block a user