wrap-langserv

#!/usr/bin/env python3

import socket
import selectors
import signal
import subprocess
import os
import sys
import traceback
import json
import time
import base64
import re

# TODO: write this in a way which would work on windows
import fcntl

from langserv import util, docstate

red = "\x1b[31m"
grn = "\x1b[32m"
yel = "\x1b[33m"
blu = "\x1b[34m"
nil = "\x1b[m"

with open("/tmp/wrap-langserv.log", "w"):
    pass

def log(*args, **kwargs):
    with open("/tmp/wrap-langserv.log", "a") as f:
        print(*args, **kwargs, file=f)


class Connection:
    """An incoming connection from outside"""
    def __init__(self, conn, injector, selector):
        self.conn = conn
        conn.setblocking(False)
        self.injector = injector
        self.selector = selector
        self.write_buf = b''
        self.parser = util.Parser(self.handle_conn_msg)

        self.selector.register(self.conn, selectors.EVENT_READ)
        self.closed = False

    def close(self):
        self.selector.unregister(self.conn)
        self.conn.close()
        self.closed = True

    def write(self, msg):
        if self.write_buf:
            self.write_buf += msg
            return
        self.write_buf = msg
        mask = selectors.EVENT_READ | selectors.EVENT_WRITE
        self.selector.modify(self.conn, mask)

    def event(self, key, mask):
        readable = mask & selectors.EVENT_READ
        writable = mask & selectors.EVENT_WRITE
        if readable:
            try:
                msg = self.conn.recv(4096)
            except ConnectionError:
                msg = None
            if not msg:
                self.close()
                return
            log(grn + msg.decode('utf8') + nil)
            self.parser.feed(msg)

        if writable:
            n = self.conn.send(self.write_buf)
            if not n:
                self.close()
                return
            self.write_buf = self.write_buf[n:]
            if not self.write_buf:
                self.selector.modify(self.conn, selectors.EVENT_READ)

    def handle_conn_msg(self, content, body, headers):
        parsed = json.loads(body)
        # injected messages must have ID's.
        if "id" not in parsed:
            raise ValueError(f"no id in {parsed}")

        def cb(_content, _parsed):
            if not self.closed():
                self.write(content)

        self.injector.inject_msg(parsed["id"], content, cb)


class LanguageServer:
    """The language server we launch as a subprocess"""
    def __init__(self, cmd, injector, selector):
        self.injector = injector
        self.selector = selector
        self.parser = util.Parser(self.handle_server_response)
        self.proc = subprocess.Popen(
            cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE
        )
        self.expect_exit = False
        self.exit_code = None
        self.write_buf = b""

        fcntl.fcntl(self.proc.stdin.fileno(), fcntl.F_SETFL, os.O_NONBLOCK)
        fcntl.fcntl(self.proc.stdout.fileno(), fcntl.F_SETFL, os.O_NONBLOCK)
        fcntl.fcntl(self.proc.stderr.fileno(), fcntl.F_SETFL, os.O_NONBLOCK)

        self.selector.register(self.proc.stdout, selectors.EVENT_READ, self)
        self.selector.register(self.proc.stderr, selectors.EVENT_READ, self)

        self.stdout_drained = False
        self.stderr_drained = False

    def close(self):
        if self.proc is None:
            return
        if not self.stdout_drained:
            self.selector.unregister(self.proc.stdout)
        if not self.stderr_drained:
            self.selector.unregister(self.proc.stderr)
        if self.write_buf:
            self.selector.unregister(self.proc.stdin)
        self.proc.terminate()
        self.exit_code = self.proc.wait()
        self.proc = None

    def write(self, msg):
        if self.proc == None:
            raise ValueError("can't write to dead server")
        if self.write_buf:
            self.write_buf += msg
            return None
        self.write_buf = msg
        self.selector.register(self.proc.stdin, selectors.EVENT_WRITE, self)

    def event(self, key, mask):
        readable = mask & selectors.EVENT_READ
        writable = mask & selectors.EVENT_WRITE
        # stdout
        if readable and key.fileobj == self.proc.stdout:
            msg = self.proc.stdout.raw.read(4096)
            if not msg:
                if not self.expect_exit:
                    raise ValueError("server died")
                self.stdout_drained = True
                self.selector.unregister(self.proc.stdout)
                return
            log(yel + msg.decode('utf8') + nil)
            self.parser.feed(msg)

        # stderr
        if readable and key.fileobj == self.proc.stderr:
            msg = self.proc.stderr.raw.read(4096)
            if not msg:
                if not self.expect_exit:
                    raise ValueError("server died")
                self.stderr_drained = True
                self.selector.unregister(self.proc.stderr)
                return
            self.injector.handle_server_error(msg)

        # stdin
        if writable:
            n = self.proc.stdin.raw.write(self.write_buf)
            if not n:
                if not self.expect_exit:
                    raise ValueError("server died")
                self.write_buf = b""
                self.selector.unregister(self.proc.stdin)
                return
            self.write_buf = self.write_buf[n:]
            if not self.write_buf:
                self.selector.unregister(self.proc.stdin)

    def handle_server_response(self, content, body, headers):
        self.injector.handle_server_response(content, body)

    def gather(self):
        if self.stdout_drained and self.stderr_drained and not self.write_buf:
            self.close()
            return True
        return False


class Editor:
    """The editor which launched us as a subprocess"""
    def __init__(self, injector, selector):
        self.injector = injector
        self.selector = selector

        fcntl.fcntl(sys.stdin.fileno(), fcntl.F_SETFL, os.O_NONBLOCK)
        fcntl.fcntl(sys.stdout.fileno(), fcntl.F_SETFL, os.O_NONBLOCK)
        self.write_buf = b""
        self.err_buf = b""
        self.selector.register(sys.stdin, selectors.EVENT_READ, self)

        self.parser = util.Parser(self.handle_stdin_msg)

    def write(self, msg):
        if self.write_buf:
            self.write_buf += msg
            return
        self.write_buf = msg
        self.selector.register(sys.stdout, selectors.EVENT_WRITE, self)

    def write_err(self, msg):
        if self.err_buf:
            self.err_buf += msg
            return
        self.err_buf = msg
        self.selector.register(sys.stderr, selectors.EVENT_WRITE, self)

    def event(self, key, mask):
        readable = mask & selectors.EVENT_READ
        writable = mask & selectors.EVENT_WRITE
        # stdin
        if readable:
            msg = os.read(sys.stdin.fileno(), 4096)
            if not msg:
                raise ValueError("stdin from editor failed")
            self.parser.feed(msg)

        # stdout
        if writable and key.fileobj == sys.stdout:
            n = os.write(sys.stdout.fileno(), self.write_buf)
            if not n:
                raise ValueError("stdout to editor failed")
            self.write_buf = self.write_buf[n:]
            if not self.write_buf:
                self.selector.unregister(sys.stdout)

        # stderr
        if writable and key.fileobj == sys.stderr:
            n = os.write(sys.stderr.fileno(), self.err_buf)
            if not n:
                raise ValueError("stderr to editor failed")
            self.err_buf = self.err_buf[n:]
            if not self.err_buf:
                self.selector.unregister(sys.stderr)

    def handle_stdin_msg(self, content, body, headers):
        # always pass the message to the language server
        self.injector.handle_editor_msg(content, body)

    def close(self):
        self.selector.unregister(sys.stdin)
        if self.write_buf:
            self.selector.unregister(sys.stdout)

    def gather(self):
        # don't exit until we have drained write_buf and err_buf
        return all([not self.write_buf, not self.err_buf])


class Timeout:
    def __init__(self, period=1.0):
        self.period = period
        self.last = time.time()

    def passed(self):
        now = time.time()
        if now > self.last + self.period:
            self.last = now
            return True
        return False


class LSPInjector:
    """Inject traffic between the Editor and LangServ based on Connections"""
    def __init__(self, cmd, selector):
        self.selector = selector
        # connections maps socket connections to Connection objects
        self.connections = {}
        # map id's to callbacks we make to handle the responses
        self.injected = {}
        # map id's to callbacks we make when we observe responses
        self.watching = {}

        if os.path.exists("/tmp/wrap-langserv.sock"):
            os.remove("/tmp/wrap-langserv.sock")

        self.listener = socket.socket(family=socket.AF_UNIX)
        self.listener.setblocking(False)
        self.listener.bind("/tmp/wrap-langserv.sock")
        self.listener.listen()

        self.ls = LanguageServer(cmd, self, selector)
        self.editor = Editor(self, selector)

        self.want_exit = False

        self.selector.register(self.listener, selectors.EVENT_READ)

        # map uri's to open Documents
        self.docs = {}

        # talon is a special Connection that we initiate and push updates to
        self.talon = None
        self.docsym = None
        self.id_gen = util.IDGen("injector")
        self.probe_talon()

    def handle_connection(self):
        conn, _ = self.listener.accept()
        self.connections[conn] = Connection(conn, self, self.selector)

    def handle_editor_msg(self, content, body):
        # write the message before we even have a chance to inject anything so
        # that text updates preceed what we inject
        self.ls.write(content)
        log(content.decode('utf8'))

        parsed = json.loads(body)
        method = parsed.get("method")
        if method == "textDocument/didOpen":
            uri = parsed["params"]["textDocument"]["uri"]
            # create a new document
            text = parsed["params"]["textDocument"]["text"]
            self.docs[uri] = docstate.Document(text)

            # request documentSymbol
            self.inject_documentSymbol(uri)

        elif method == "textDocument/didClose":
            uri = parsed["params"]["textDocument"]["uri"]
            # discard document state
            del self.docs[uri]

        elif method == "textDocument/didChange":
            uri = parsed["params"]["textDocument"]["uri"]
            # sync changes
            doc = self.docs[uri]
            for change in parsed["params"]["contentChanges"]:
                doc.did_change(docstate.ContentChange(change))

            # request documentSymbol
            self.inject_documentSymbol(uri)

        elif method == "textDocument/completion":
            uri = parsed["params"]["textDocument"]["uri"]
            id = parsed["id"]
            pos = docstate.Position(parsed["params"]["position"])
            self.watch_completion(uri, id, pos)

        elif method == "exit":
            self.ls.expect_exit = True
            self.want_exit = True

    def inject_documentSymbol(self, uri):
        id = next(self.id_gen)
        msg = {
            "jsonrpc": "2.0",
            "method": "textDocument/documentSymbol",
            "id": id,
            "params": {
                "textDocument": {
                    "uri": uri
                }
            }
        }

        content = util.make_content(msg)

        def docsym_cb(_, parsed):
            headers = {"Uri": uri, "Type": "documentSymbol"}
            _content = util.make_content(parsed, headers)
            self.docsym = _content
            if self.talon is not None:
                self.talon.write(_content)

        self.inject_msg(id, content, docsym_cb)

    def watch_completion(self, uri, id, pos):
        text = self.docs[uri].text
        # give talon the line preceeding the completion request, so it can
        # figure out how much of the word to not type
        pretext = ""
        if pos.character > 0:
            line_start = docstate.Position(line=pos.line, character=0)
            start = line_start.index(text)
            end = pos.index(text)
            pretext = text[start:end]
        pretext = base64.b64encode(pretext.encode('utf8')).decode('utf8')

        def cb(_, _parsed):
            headers = {
                "Uri": uri,
                "Type": "completion",
                "Pretext": pretext,
            }
            _content = util.make_content(_parsed, headers)
            if self.talon is not None:
                self.talon.write(_content)

        self.watch_msg(id, cb)

    def watch_msg(self, id, cb):
        self.watching[id] = cb

    def inject_msg(self, id, content, cb):
        self.injected[id] = cb
        self.ls.write(content)
        log(blu + content.decode('utf8') + nil)

    def handle_server_response(self, content, body):
        parsed = json.loads(body)
        # do we need to direct this response to a connection?
        if "id" in parsed and parsed["id"] in self.injected:
            cb = self.injected.pop(parsed["id"])
            cb(content, parsed)
            return

        self.editor.write(content)

        # did we want that message for something?
        if "id" in parsed and parsed["id"] in self.watching:
            cb = self.watching.pop(parsed["id"])
            cb(content, parsed)

    def handle_server_error(self, msg):
        # log the message in red
        log(red + msg.decode('utf8') + nil, end="")
        # pass the message to the editor.
        self.editor.write_err(msg)

    def close(self):
        for conn in self.connections.values():
            conn.close()
        self.ls.close()
        self.editor.close()
        self.selector.unregister(self.listener)
        self.listener.close()

    def __enter__(self):
        return self

    def __exit__(self, etype, value, tb):
        self.close()

    def run(self):
        timeout = Timeout(period=1.0)
        while not self.gather():
            for key, mask in self.selector.select(timeout.period):
                readable = mask & selectors.EVENT_READ
                writable = mask & selectors.EVENT_WRITE
                if key.data == self.editor:
                    self.editor.event(key, mask)
                elif key.data == self.ls:
                    self.ls.event(key, mask)
                elif key.fileobj == self.listener:
                    self.handle_connection()
                elif self.talon and key.fileobj == self.talon.conn:
                    self.talon.event(key, mask)
                    if self.talon.closed:
                        self.talon = None
                else:
                    conn = self.connections[key.fileobj]
                    conn.event(key, mask)
                    if conn.closed:
                        del self.connections[key.fileobj]
            if timeout.passed():
                self.probe_talon()

    def gather(self):
        # see if we are ready to exit
        return all([
            self.want_exit,
            self.editor.gather(),
            self.ls.gather(),
        ])

    def probe_talon(self):
        if self.talon is not None:
            return

        sockpath = "~/.talon/user/rexroni_talon/langserv/langserv.sock"
        sockpath = os.path.expanduser(sockpath)

        sock = socket.socket(family=socket.AF_UNIX)
        try:
            sock.connect(sockpath)
        except (ConnectionError, FileNotFoundError):
            sock.close()
            return

        log('connected to talon!')

        self.talon = Connection(sock, self, self.selector)

        if self.docsym is not None:
            self.talon.write(self.docsym)


def main(cmd):
    exit_code = 0
    try:
        with selectors.DefaultSelector() as selector:
            with LSPInjector(cmd, selector) as injector:
                injector.run()

            if injector.ls.exit_code is not None:
                exit_code = injector.ls.exit_code

    except Exception as e:
        with open("/tmp/wrap-langserv.log", "a") as f:
            traceback.print_exc(file=f)
        # We're out of the event loop, but still give the editor a one-liner
        print(e, file=sys.stderr)
        exit(1)

    exit(exit_code)

if __name__ == "__main__":
    # check for path hack
    executable = os.path.basename(sys.argv[0])
    if executable != "wrap-langserv":
        # insert the name of the executable we were called with before the args
        args = [executable] + sys.argv[1:]
    else:
        args = sys.argv[1:]

    # regardless of the path hack, always remove the path hack before calling
    # the subprocess, in case the path hack is active but we were invoked
    # directly
    path = os.environ.get("PATH")
    if path is not None:
        path = re.sub("[^:]*/langserv/path-hack/?:", "", path)
        os.environ["PATH"] = path

    main(args)