Berry/lsp_client.py

import asyncio
import json
import time
import lsprotocol

from pathlib import Path


class NoFileOpen(Exception):
    def __init__(self, *args):
        super().__init__(self, *args)

class LSPClient:
    def __init__(self, proc_name: str, proc_flags: list[str], langauge_id: str):
        self.proc = None
        self.proc_flags = proc_flags
        self.proc_name = proc_name
        self.request_id = 0
        
        self.langauge_id = langauge_id
        self.file_path = None
        self.folder = None
        
        self.has_file_open = False
        
        self.pending = {}
        self.running = False
        
    async def message_loop(self):
        while True:
            msg = await self.read_message()
            
            if "id" in msg:
                future = self.pending.pop(msg["id"], None)
                
                if future:
                    future.set_result(msg)
            elif "method" in msg:
                print("notification:", msg)
        
    async def open_folder(self, folder_path: str):
        self.folder = Path(folder_path).resolve().as_uri()
        
    async def open_file(self, file_path: str):
        
        
        with open(file_path, "r") as f:
            content = f.read()
            
        resolved_path = Path(file_path).resolve()
        self.file_path = resolved_path.as_uri()
        await self.open_folder(resolved_path.parent)

        await self.send({
            "jsonrpc": "2.0",
            "method": "textDocument/didOpen",
            "params": {
                "textDocument": {
                    "uri": self.file_path,
                    "languageId": self.langauge_id,
                    "version": 1,
                    "text": content
                }
            }
        })
        
        self.has_file_open = True
        
    async def start(self, starting_file_path: str):
        if self.running:
            raise Exception("LSP is already running!")
        self.running = True
        
        resolved_path = Path(starting_file_path).resolve()
        
        self.file_path = resolved_path.as_uri()
        self.folder = resolved_path.parent.as_uri()
        
        # create a subprocess which we talk to over stdio
        self.proc = await asyncio.create_subprocess_exec(
            self.proc_name,
            *self.proc_flags,
            stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE
        )
        
        asyncio.create_task(self.message_loop())
        
        await self.initialize()
        await self.open_file(starting_file_path)
        
    async def send(self, data: dict):
        body = json.dumps(data).encode("utf-8")
        
        header = (
            f"Content-Length: {len(body)}\r\n\r\n"
        ).encode("ascii")
        
        print("writing")
        self.proc.stdin.write(header + body)

        print("draining")
        await self.proc.stdin.drain()
        print("done draining")
        
    async def read_message(self):
        headers = {}
        
        while True:
            line = await self.proc.stdout.readline()
            
            if line == b"\r\n": break
            
            decoded = line.decode("ascii").strip()
            
            key, value = decoded.split(":", 1)
            headers[key.strip()] = value.strip()
            
        content_length = int(headers["Content-Length"])
        
        body = await self.proc.stdout.readexactly(content_length)
        
        return json.loads(body.decode("utf-8"))
        
    async def apply_change(self, text: str):
        await self.send({
            "jsonrpc": "2.0",
            "method": "textDocument/didChange",
            "params": {
                "textDocument": {
                    "uri": self.file_path,
                    "version": 1
                },
                "contentChanges": [
                    {
                        "text": text
                    }
                ]
            }
        })
        
    async def request(self, method, params):
        self.request_id += 1
        reg_id = self.request_id
        
        loop = asyncio.get_running_loop()
        
        future = loop.create_future()
        
        print(self.request_id)
        
        self.pending[reg_id] = future
        print(self.pending)
        
        await self.send({
            "jsonrpc": "2.0",
            "id": self.request_id,
            "method": method,
            "params": params
        })
        
        return await future

    async def get_completions(self, cursor_pos: tuple[int, int]):
        if not self.has_file_open:
            raise NoFileOpen("Can't get completions when you haven't openned a file yet.")
        
        return await self.request(
            "textDocument/completion",
            {
                "textDocument": {
                    "uri": self.file_path,
                },
                "position": {
                    "line": cursor_pos[0],
                    "character": cursor_pos[1]
                },
                "context": {
                    "triggerKind": 1
                }
            }
        )
        
    async def initialize(self):
        print(f"rootUri: {self.folder}")
        
        response = await self.request(
            "initialize",
            {
                "processId": None,
                "rootUri": self.folder,
                "capabilities": {},
                "initializationOptions": {
                    "clangdFileStatus": True
                }
            }
        )
        
        print(response)
        
        await self.send({
            "jsonrpc": "2.0",
            "method": "initialized",
            "params": {}
        })
        
    def get_from_file_path(file_path: str, ):
        file_path = Path(file_path)
        
        
        
        
if __name__ == "__main__":
    async def main():
        client = LSPClient("jedi-language-server", [], "python")
        await client.start("test.py")
        
        start = time.time()
        print(await client.get_completions((0, 9)))
        
        
        
        print(f"took {round(time.time() - start, 1)} seconds")
        
    asyncio.run(main())