nx_post_support/server/src/lsp_server.py

# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
"""Implementation of tool support over LSP."""

from __future__ import annotations

import copy
import json
import os
import pathlib
import re
import sys
import sysconfig
import traceback
from typing import Any, List, Optional, Sequence, Tuple
import re


# **********************************************************
# Update sys.path before importing any bundled libraries.
# **********************************************************
def update_sys_path(path_to_add: str, strategy: str) -> None:
    """Add given path to `sys.path`."""
    if path_to_add not in sys.path and os.path.isdir(path_to_add):
        if strategy == "useBundled":
            sys.path.insert(0, path_to_add)
        elif strategy == "fromEnvironment":
            sys.path.append(path_to_add)


# Ensure that we can import LSP libraries, and other bundled libraries.
update_sys_path(
    os.fspath(pathlib.Path(__file__).parent.parent / "libs"),
    os.getenv("LS_IMPORT_STRATEGY", "useBundled"),
)

# **********************************************************
# Imports needed for the language server goes below this.
# **********************************************************
# pylint: disable=wrong-import-position,import-error
import lsp_jsonrpc as jsonrpc
import lsp_utils as utils
import lsprotocol.types as lsp
from pygls import server, uris, workspace
from pygls.workspace.text_document import TextDocument
from common.load_data import standard_items
from common.formatter import format_tcl
from tclint.parser import Parser
from tclint.lexer import TclSyntaxError
from tclint.format import Formatter, FormatterOpts
from tclint.violations import Violation
from tools.semantic_tokens import (
    SemanticTokenCollector,
    collect_semantic_tokens,
    encode_tokens,
)
from nx_plugins.poco_plugin import commands
from tools import poco_check

DIAGNOSTIC_SOURCE = "tclint"


class TclLanguageServer(server.LanguageServer):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.parser = Parser()
        self.parser._commands.update(commands)
        self.diagnostics = {}

    def format(
        self,
        document: TextDocument,
        options: lsp.FormattingOptions,
        range: Optional[Tuple[int, int]] = None,
    ):
        # parser = Parser(command_plugins=["nx_plugins.poco_plugin.py"])
        # parser._commands.update(commands)

        indent = "\t" if not options.insert_spaces else " " * options.tab_size

        formatter = Formatter(
            FormatterOpts(
                indent=indent,
                spaces_in_braces=False,
                max_blank_lines=500,
                indent_namespace_eval=True,
            ),
        )

        if range is not None:
            start, end = range
            return formatter.format_partial(document.source[start:end], self.parser)

        return formatter.format_top(document.source, self.parser)

    def linter(
        self,
        document: TextDocument,
    ) -> List[Violation]:
        violations = []
        tree = self.parser.parse(document.source)
        violations += self.parser.violations

        # if debug > 0:
        #     print(tree.pretty(positions=(debug > 1)))

        for checker in poco_check.get_checkers():
            violations += checker.check(document.source, tree)

        # v = CommentVisitor()
        # ignore_lines = v.run(tree, path)
        # violations = filter_violations(violations, config.ignore, ignore_lines)

        return violations

    def lint(self, document: TextDocument):
        diagnostics = []

        try:
            violations = self.linter(document)
        except TclSyntaxError as e:
            return [
                lsp.Diagnostic(
                    message=str(e),
                    severity=lsp.DiagnosticSeverity.Error,
                    range=lsp.Range(
                        start=lsp.Position(e.start[0] - 1, e.start[1] - 1),
                        end=lsp.Position(e.end[0] - 1, e.end[1] - 1),
                    ),
                    code="syntax error",
                    source=DIAGNOSTIC_SOURCE,
                )
            ]

        for violation in violations:
            message = violation.message
            severity = lsp.DiagnosticSeverity.Warning
            start = lsp.Position(
                line=violation.start[0] - 1, character=violation.start[1] - 1
            )
            end = lsp.Position(
                line=violation.end[0] - 1, character=violation.end[1] - 1
            )

            diagnostics.append(
                lsp.Diagnostic(
                    message=message,
                    severity=severity,
                    range=lsp.Range(
                        start=start,
                        end=end,
                    ),
                    code=violation.id,
                    source=DIAGNOSTIC_SOURCE,
                )
            )

        return diagnostics

    def _compute_diagnostics(self, document: TextDocument) -> List[lsp.Diagnostic]:
        return self.lint(document)

    def compute_diagnostics(self, document: TextDocument):
        # `None` sentinel ensures that `diagnostics` gets updated if the URI is not
        # present.
        _, previous = self.diagnostics.get(document, (0, None))

        diagnostics = self._compute_diagnostics(document)

        # Only update if the list has changed
        if previous != diagnostics:
            self.diagnostics[document.uri] = (document.version, diagnostics)


WORKSPACE_SETTINGS = {}
GLOBAL_SETTINGS = {}
RUNNER = pathlib.Path(__file__).parent / "lsp_runner.py"

TOKEN_TYPES = [
    "command",
    "variable",
    "function",
    "string",
    "number",
    "keyword",
    "comment",
]
TOKEN_MODIFIERS = []

MAX_WORKERS = 5
LSP_SERVER = TclLanguageServer(
    name="NX Postprocessor Support", version="0.0.1", max_workers=MAX_WORKERS
)

# **********************************************************
# Tool specific code goes below this.
# **********************************************************
TOOL_MODULE = "nx-post-support"
TOOL_DISPLAY = "NX Postprocessor Support"
TOOL_ARGS = []  # default arguments always passed to your tool.

# Delete "Linting features" section if your tool is NOT a linter.
# **********************************************************
# Linting features start here
# **********************************************************

#  See `pylint` implementation for a full featured linter extension:
#  Pylint: https://github.com/microsoft/vscode-pylint/blob/main/bundled/tool


@LSP_SERVER.feature(lsp.TEXT_DOCUMENT_DID_OPEN)
def did_open(params: lsp.DidOpenTextDocumentParams) -> None:
    """LSP handler for textDocument/didOpen request."""
    document = LSP_SERVER.workspace.get_text_document(params.text_document.uri)
    LSP_SERVER.compute_diagnostics(document)


@LSP_SERVER.feature(lsp.TEXT_DOCUMENT_DID_SAVE)
def did_save(params: lsp.DidSaveTextDocumentParams) -> None:
    """LSP handler for textDocument/didSave request."""
    document = LSP_SERVER.workspace.get_text_document(params.text_document.uri)


@LSP_SERVER.feature(lsp.TEXT_DOCUMENT_DID_CLOSE)
def did_close(params: lsp.DidCloseTextDocumentParams) -> None:
    """LSP handler for textDocument/didClose request."""


@LSP_SERVER.feature(lsp.TEXT_DOCUMENT_DID_CHANGE)
def did_change(params: lsp.DidChangeTextDocumentParams) -> None:
    """LSP handler for textDocument/didChange request"""
    document = LSP_SERVER.workspace.get_text_document(params.text_document.uri)
    LSP_SERVER.compute_diagnostics(document)


@LSP_SERVER.feature(
    lsp.TEXT_DOCUMENT_DIAGNOSTIC,
    lsp.DiagnosticOptions(
        identifier="pull-diagnostics",
        inter_file_dependencies=False,
        workspace_diagnostics=False,
    ),
)
def document_diagnostic(params: lsp.DocumentDiagnosticParams):
    """Return diagnostics for the requested document"""
    was_cached = True
    if (uri := params.text_document.uri) not in LSP_SERVER.diagnostics:
        was_cached = False
        doc = LSP_SERVER.workspace.get_text_document(uri)
        LSP_SERVER.compute_diagnostics(doc)

    version, diagnostics = LSP_SERVER.diagnostics[uri]
    result_id = f"{uri}@{version}"

    if was_cached and result_id == params.previous_result_id:
        return lsp.UnchangedDocumentDiagnosticReport(result_id)

    return lsp.FullDocumentDiagnosticReport(items=diagnostics, result_id=result_id)


@LSP_SERVER.feature(lsp.TEXT_DOCUMENT_COMPLETION)
def on_completion(params: lsp.CompletionParams) -> list[lsp.CompletionItem]:
    items = (
        standard_items.tcl_keyword_list
        + standard_items.nx_procs
        + standard_items.nx_variables
    )
    return lsp.CompletionList(is_incomplete=False, items=items)


@LSP_SERVER.feature(lsp.TEXT_DOCUMENT_SEMANTIC_TOKENS_FULL)
def on_semantic_tokens(params: lsp.SemanticTokensParams):
    doc = LSP_SERVER.workspace.get_document(params.text_document.uri)
    code = doc.source

    tokens = collect_semantic_tokens(code)
    data = encode_tokens(tokens)

    return lsp.SemanticTokens(data=data)


@LSP_SERVER.feature(lsp.TEXT_DOCUMENT_HOVER)
def hover(params: lsp.HoverParams) -> lsp.Hover:
    pos = params.position
    document_uri = params.text_document.uri
    document = LSP_SERVER.workspace.get_text_document(document_uri)
    col = params.position.character

    try:
        line = document.lines[pos.line]
    except IndexError:
        return None

    for m in re.finditer(r"\b\w+\b", line):
        if m.start() <= col <= m.end():
            token = m.group(0)
            break
    else:
        return None

    command = token
    data = standard_items.json_data

    all_items = data.get("MOM_procs", []) + data.get("mom_variables", [])

    match = next((item for item in all_items if item["label"] == command), None)
    if not match:
        return None

    if not match.get("kind") == "function":
        return None

    label = match.get("label", "")

    parameters = match.get("parameters", [])
    param_lines = (
        "\n".join(f"- `{p['name']}`: {p['desc']}" for p in parameters) or "_None_"
    )

    example_data = match.get("example", [])
    example_md = "\n".join(f"{line}" for line in example_data)

    returns_data = match.get("returns", ["None"])
    returns_md = "\n".join(f"- {line}" for line in returns_data)

    doc_md = f"""\
### 📘 {label}

**Purpose**  
{match.get("description", "No description available.")}

**Format**  
`{match.get("format", label)}`

**Parameters**  
{param_lines}

**Return value**  
{returns_md}

**Example**  
```tcl
{example_md}"""

    return lsp.Hover(lsp.MarkupContent(kind=lsp.MarkupKind.Markdown, value=doc_md))


# **********************************************************
# Linting features end here
# **********************************************************


# **********************************************************
# Formatting features start here
# **********************************************************
#  Sample implementations:
#  Black: https://github.com/microsoft/vscode-black-formatter/blob/main/bundled/tool


# **********************************************************
# Formatting features ends here
# **********************************************************
@LSP_SERVER.feature(lsp.TEXT_DOCUMENT_FORMATTING)
def formatting(params: lsp.DocumentFormattingParams) -> list[lsp.TextEdit] | None:
    """LSP handler for textDocument/formatting request."""
    doc = LSP_SERVER.workspace.get_text_document(params.text_document.uri)

    source = doc.source
    start = lsp.Position(line=0, character=0)
    last_line = source.rsplit("\n", 1)[-1]
    end = lsp.Position(line=source.count("\n"), character=len(last_line))

    formatted = LSP_SERVER.format(doc, params.options)
    return [
        lsp.TextEdit(
            range=lsp.Range(start=start, end=end),
            new_text=formatted,
        )
    ]


# **********************************************************
# Required Language Server Initialization and Exit handlers.
# **********************************************************
@LSP_SERVER.feature(lsp.WORKSPACE_DID_CHANGE_CONFIGURATION)
def did_change_configuration(params: lsp.DidChangeConfigurationParams):
    """LSP Handler for Config Changes"""


@LSP_SERVER.feature(lsp.INITIALIZE)
def initialize(params: lsp.InitializeParams) -> lsp.InitializeResult:
    """LSP handler for initialize request."""
    log_to_output(f"CWD Server: {os.getcwd()}")

    paths = "\r\n   ".join(sys.path)
    log_to_output(f"sys.path used to run Server:\r\n   {paths}")

    GLOBAL_SETTINGS.update(**params.initialization_options.get("globalSettings", {}))

    settings = params.initialization_options["settings"]
    _update_workspace_settings(settings)
    log_to_output(
        f"Settings used to run Server:\r\n{json.dumps(settings, indent=4, ensure_ascii=False)}\r\n"
    )
    log_to_output(
        f"Global settings:\r\n{json.dumps(GLOBAL_SETTINGS, indent=4, ensure_ascii=False)}\r\n"
    )
    semantic_tokens_legend = lsp.SemanticTokensLegend(
        token_types=TOKEN_TYPES,
        token_modifiers=TOKEN_MODIFIERS,
    )
    return lsp.InitializeResult(
        capabilities=lsp.ServerCapabilities(
            document_formatting_provider=GLOBAL_SETTINGS.get("formatter", True),
            semantic_tokens_provider=lsp.SemanticTokensOptions(
                legend=semantic_tokens_legend, full=True, range=False
            ),
        )
    )


@LSP_SERVER.feature(lsp.EXIT)
def on_exit(_params: Optional[Any] = None) -> None:
    """Handle clean up on exit."""
    jsonrpc.shutdown_json_rpc()


@LSP_SERVER.feature(lsp.SHUTDOWN)
def on_shutdown(_params: Optional[Any] = None) -> None:
    """Handle clean up on shutdown."""
    jsonrpc.shutdown_json_rpc()


def _get_global_defaults():
    return {
        "path": GLOBAL_SETTINGS.get("path", []),
        "interpreter": GLOBAL_SETTINGS.get("interpreter", [sys.executable]),
        "args": GLOBAL_SETTINGS.get("args", []),
        "importStrategy": GLOBAL_SETTINGS.get("importStrategy", "useBundled"),
        "showNotifications": GLOBAL_SETTINGS.get("showNotifications", "off"),
        "formatter": GLOBAL_SETTINGS.get("formatter", True),
    }


def _update_workspace_settings(settings):
    if not settings:
        key = os.getcwd()
        WORKSPACE_SETTINGS[key] = {
            "cwd": key,
            "workspaceFS": key,
            "workspace": uris.from_fs_path(key),
            **_get_global_defaults(),
        }
        return

    for setting in settings:
        key = uris.to_fs_path(setting["workspace"])
        WORKSPACE_SETTINGS[key] = {
            "cwd": key,
            **setting,
            "workspaceFS": key,
        }


def _get_settings_by_path(file_path: pathlib.Path):
    workspaces = {s["workspaceFS"] for s in WORKSPACE_SETTINGS.values()}

    while file_path != file_path.parent:
        str_file_path = str(file_path)
        if str_file_path in workspaces:
            return WORKSPACE_SETTINGS[str_file_path]
        file_path = file_path.parent

    setting_values = list(WORKSPACE_SETTINGS.values())
    return setting_values[0]


def _get_document_key(document: workspace.Document):
    if WORKSPACE_SETTINGS:
        document_workspace = pathlib.Path(document.path)
        workspaces = {s["workspaceFS"] for s in WORKSPACE_SETTINGS.values()}

        # Find workspace settings for the given file.
        while document_workspace != document_workspace.parent:
            if str(document_workspace) in workspaces:
                return str(document_workspace)
            document_workspace = document_workspace.parent

    return None


def _get_settings_by_document(document: workspace.Document | None):
    if document is None or document.path is None:
        return list(WORKSPACE_SETTINGS.values())[0]

    key = _get_document_key(document)
    if key is None:
        # This is either a non-workspace file or there is no workspace.
        key = os.fspath(pathlib.Path(document.path).parent)
        return {
            "cwd": key,
            "workspaceFS": key,
            "workspace": uris.from_fs_path(key),
            **_get_global_defaults(),
        }

    return WORKSPACE_SETTINGS[str(key)]


# *****************************************************
# Logging and notification.
# *****************************************************
def log_to_output(
    message: str, msg_type: lsp.MessageType = lsp.MessageType.Log
) -> None:
    LSP_SERVER.show_message_log(message, msg_type)


def log_error(message: str) -> None:
    LSP_SERVER.show_message_log(message, lsp.MessageType.Error)
    if os.getenv("LS_SHOW_NOTIFICATION", "off") in ["onError", "onWarning", "always"]:
        LSP_SERVER.show_message(message, lsp.MessageType.Error)


def log_warning(message: str) -> None:
    LSP_SERVER.show_message_log(message, lsp.MessageType.Warning)
    if os.getenv("LS_SHOW_NOTIFICATION", "off") in ["onWarning", "always"]:
        LSP_SERVER.show_message(message, lsp.MessageType.Warning)


def log_always(message: str) -> None:
    LSP_SERVER.show_message_log(message, lsp.MessageType.Info)
    if os.getenv("LS_SHOW_NOTIFICATION", "off") in ["always"]:
        LSP_SERVER.show_message(message, lsp.MessageType.Info)


# *****************************************************
# Start the server.
# *****************************************************
if __name__ == "__main__":
    LSP_SERVER.start_io()