A Vyper LSP
import sys
from pygls.lsp.types.language_features import List
from vyper.ast import VyperNode, nodes
from vyper.compiler import CompilerData

ast = None


class AST:
    _instance = None
    ast_data = None

    custom_type_node_types = (
        nodes.StructDef,
        nodes.EnumDef,
        nodes.EventDef
    )

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super(AST, cls).__new__(cls)
            cls._instance.ast_data = None
        return cls._instance

    def update_ast(self, document):
        try:
            compiler_data = CompilerData(document.source)
            self.ast_data = compiler_data.vyper_module_unfolded
        except Exception as e:
            print(f"Error updating AST: {e}", file=sys.stderr)
            pass

    def build_ast(self, src: str):
        try:
            compiler_data = CompilerData(src)
            self.ast_data = compiler_data.vyper_module_unfolded
        except Exception:
            pass

    def get_enums(self) -> List[str]:
        if self.ast_data is None:
            return []

        return [
            node.name
            for node in self.ast_data.get_descendants(nodes.EnumDef)
        ]

    def get_structs(self) -> List[str]:
        if self.ast_data is None:
            return []

        return [
            node.name
            for node in self.ast_data.get_descendants(nodes.StructDef)
        ]

    def get_events(self) -> List[str]:
        if self.ast_data is None:
            return []

        return [
            node.name
            for node in self.ast_data.get_descendants(nodes.EventDef)
        ]

    def get_user_defined_types(self):
        if self.ast_data is None:
            return []

        return [
            node.name
            for node in self.ast_data.get_descendants(self.custom_type_node_types)
        ]

    def get_constants(self):
        if self.ast_data is None:
            return []

        return [
            node.target.id
            for node in self.ast_data.get_descendants(nodes.VariableDecl)
            if node.is_constant
        ]


    def get_enum_variants(self, enum: str):
        if self.ast_data is None:
            return []

        enum_node = self.find_type_declaration_node(enum)
        if enum_node is None:
            return []

        return [node.value.id for node in enum_node.get_children()]

    def get_state_variables(self):
        if self.ast_data is None:
            return []

        return [
            node.target.id
            for node in self.ast_data.get_descendants(nodes.VariableDecl)
        ]

    def get_internal_functions(self):
        if self.ast_data is None:
            return []

        function_nodes = self.ast_data.get_descendants(nodes.FunctionDef)
        inernal_nodes = []

        for node in function_nodes:
            for decorator in node.decorator_list:
                if decorator.id == "internal":
                    inernal_nodes.append(node)

        return inernal_nodes

    def find_nodes_referencing_internal_function(self, function: str):
        if self.ast_data is None:
            return []

        return self.ast_data.get_descendants(nodes.Call, {"func.attr": function, "func.value.id": "self"})


    def find_nodes_referencing_state_variable(self, variable: str):
        if self.ast_data is None:
            return []

        return self.ast_data.get_descendants(nodes.Attribute, {"value.id": "self", "attr": variable})

    def find_nodes_referencing_constant(self, constant: str):
        if self.ast_data is None:
            return []

        name_nodes = self.ast_data.get_descendants(nodes.Name, {"id": constant})
        return [node for node in name_nodes if not isinstance(node.get_ancestor(), nodes.VariableDecl)]

    def get_attributes_for_symbol(self, symbol: str):
        if self.ast_data is None:
            return []

        node = self.find_type_declaration_node(symbol)
        if node is None:
            return []

        if isinstance(node, nodes.StructDef):
            return []
        elif isinstance(node, nodes.EnumDef):
            return self.get_enum_variants(symbol)
        else:
            return []

    def find_function_declaration_node_for_name(self, function: str):
        if self.ast_data is None:
            return None

        for node in self.ast_data.get_descendants(nodes.FunctionDef):
            name_match = node.name == function
            not_interface_declaration = not isinstance(node.get_ancestor(), nodes.InterfaceDef)
            if name_match and not_interface_declaration:
                return node

        return None


    def find_state_variable_declaration_node(self, variable: str):
        if self.ast_data is None:
            return None

        for node in self.ast_data.get_descendants(nodes.VariableDecl):
            if node.target.id == variable:
                return node

        return None

    def find_type_declaration_node(self, symbol: str):
        if self.ast_data is None:
            return None

        for node in self.ast_data.get_descendants(self.custom_type_node_types):
            if node.name == symbol:
                return node
            if isinstance(node, nodes.EnumDef):
                for variant in node.get_children(nodes.Expr):
                    if variant.value.id == symbol:
                        return variant

        return None

    def find_nodes_referencing_enum(self, enum: str):
        if self.ast_data is None:
            return []

        return_nodes = []

        for node in self.ast_data.get_descendants(nodes.AnnAssign, {"annotation.id": enum}):
            return_nodes.append(node)
        for node in self.ast_data.get_descendants(nodes.Attribute, {"value.id": enum}):
            return_nodes.append(node)
        for node in self.ast_data.get_descendants(nodes.VariableDecl, {"annotation.id": enum}):
            return_nodes.append(node)

        return return_nodes

    def find_nodes_referencing_enum_variant(self, enum: str, variant: str):
        if self.ast_data is None:
            return None

        return self.ast_data.get_descendants(nodes.Attribute, {"attr": variant, "value.id": enum})

    def find_nodes_referencing_struct(self, struct: str):
        if self.ast_data is None:
            return []

        return_nodes = []

        for node in self.ast_data.get_descendants(nodes.AnnAssign, {"annotation.id": struct}):
            return_nodes.append(node)
        for node in self.ast_data.get_descendants(nodes.Call, {"func.id": struct}):
            return_nodes.append(node)
        for node in self.ast_data.get_descendants(nodes.VariableDecl, {"annotation.id": struct}):
            return_nodes.append(node)
        for node in self.ast_data.get_descendants(nodes.FunctionDef, {"returns.id": struct}):
            return_nodes.append(node)

        return return_nodes

src/ast.py
vyperlang/vyper-lsp
import sys

from pygls.lsp.types.language_features import (
    CompletionItem,
    CompletionList,
    CompletionParams,
)
from pygls.server import LanguageServer

from src.ast import AST
from src.types import BASE_TYPES

DECORATORS = ["payable", "nonpayable", "view", "pure", "external", "internal"]


class Completer:
    def __init__(self, ast=None):
        self.ast = ast or AST()

    def get_completions(
        self, ls: LanguageServer, params: CompletionParams
    ) -> CompletionList:
        items = []
        document = ls.workspace.get_document(params.text_document.uri)
        og_line = document.lines[params.position.line]
        current_line = document.lines[params.position.line].strip()
        custom_types = self.ast.get_user_defined_types()

        if params.context:
            if params.context.trigger_character == ".":
                # get element before the dot
                element = current_line.split(" ")[-1].split(".")[0]
                for attr in self.ast.get_attributes_for_symbol(element):
                    items.append(CompletionItem(label=attr))
                l = CompletionList(is_incomplete=False, items=[])
                l.add_items(items)
                return l
            elif params.context.trigger_character == "@":
                for dec in DECORATORS:
                    items.append(CompletionItem(label=dec))
                l = CompletionList(is_incomplete=False, items=[])
                l.add_items(items)
                return l
            elif params.context.trigger_character == ":":
                for typ in custom_types + list(BASE_TYPES):
                    items.append(CompletionItem(label=typ, insert_text=f" {typ}"))

                l = CompletionList(is_incomplete=False, items=[])
                l.add_items(items)
                return l
            else:
                if params.context.trigger_character == " ":
                    if current_line[-1] == ":":
                        print("triggered type", file=sys.stderr)
                        for typ in custom_types + list(BASE_TYPES):
                            items.append(CompletionItem(label=typ))

                        l = CompletionList(is_incomplete=False, items=[])
                        l.add_items(items)
                        return l

                print(
                    f"unknown trigger character {params.context.trigger_character}",
                    file=sys.stderr,
                )
        else:
            print("no context", file=sys.stderr)
            print(params, file=sys.stderr)
        return CompletionList(is_incomplete=False, items=[])

src/completions.py
vyperlang/vyper-lsp
import sys
import warnings
from lark import UnexpectedInput, UnexpectedToken
from pygls.lsp.types.language_features import Diagnostic, List, Position, Range
from pygls.workspace import Document
from vyper.compiler import CompilerData
from vyper.exceptions import VyperException

from .parse import parser
import re

pattern = r"(.+) is deprecated\. Please use `(.+)` instead\."
compiled_pattern = re.compile(pattern)

def format_parse_error(e):
    if isinstance(e, UnexpectedToken):
        expected = ", ".join(e.accepts or e.expected)
        return f"Unexpected token '{e.token}' at {e.line}:{e.column}. Expected one of: {expected}"
    else:
        return str(e)


def get_diagnostics(doc: Document):
    doctext = doc.source
    diagnostics: List[Diagnostic] = []
    last_error = None

    def on_error(e: UnexpectedInput):
        nonlocal last_error
        if (
            last_error is not None
            and last_error.line == e.line
            and last_error.column == e.column
        ):
            return
        if (
            last_error is not None
            and type(last_error) == type(e)
            and last_error.line == e.line
        ):
            return
        diagnostics.append(
            Diagnostic(
                range=Range(
                    start=Position(line=e.line - 1, character=e.column - 1),
                    end=Position(line=e.line - 1, character=e.column),
                ),
                message=format_parse_error(e),
            )
        )
        last_error = e
        return True

    try:
        parser.parse(doctext, on_error=on_error)
    except Exception:
        # ignore errors that are already handled by on_error
        pass

    compiler_data = CompilerData(doctext)

    warnings.simplefilter("always")
    replacements = {}
    with warnings.catch_warnings(record=True) as w:
        try:
            compiler_data.vyper_module_unfolded
        except VyperException as e:
            if e.lineno is not None and e.col_offset is not None:
                diagnostics.append(
                    Diagnostic(
                        range=Range(
                            start=Position(line=e.lineno - 1, character=e.col_offset - 1),
                            end=Position(line=e.lineno - 1, character=e.col_offset),
                        ),
                        message=str(e),
                        severity=1,
                    )
                )
            else:
                for a in e.annotations:
                    diagnostics.append(
                        Diagnostic(
                            range=Range(
                                start=Position(
                                    line=a.lineno - 1, character=a.col_offset - 1
                                ),
                                end=Position(line=a.lineno - 1, character=a.col_offset),
                            ),
                            message=e.message,
                            severity=1,
                        )
                    )
        print(f"{len(w)} warnings", file=sys.stderr)
        for warning in w:
            match = compiled_pattern.match(str(warning.message))
            if not match:
                continue
            deprecated = match.group(1)
            replacement = match.group(2)
            replacements[deprecated] = replacement

    # iterate over doc.lines and find all deprecated values
    # and create a warning for each one at the correct position
    for i, line in enumerate(doc.lines):
        for deprecated, replacement in replacements.items():
            if deprecated in line:
                diagnostics.append(
                    Diagnostic(
                        range=Range(
                            start=Position(line=i, character=line.index(deprecated)),
                            end=Position(line=i, character=line.index(deprecated) + len(deprecated)),
                        ),
                        message=f"{deprecated} is deprecated. Please use {replacement} instead.",
                        severity=2,
                    )
                )

    return diagnostics

src/diagnostics.py
vyperlang/vyper-lsp
from typing import Optional
from lsprotocol.types import (
    TEXT_DOCUMENT_COMPLETION,
    TEXT_DOCUMENT_DID_CHANGE,
    TEXT_DOCUMENT_DID_OPEN,
    TEXT_DOCUMENT_DECLARATION,
    TEXT_DOCUMENT_DEFINITION,
    TEXT_DOCUMENT_REFERENCES,
)
from pygls.lsp.methods import HOVER, IMPLEMENTATION
from pygls.lsp.types import (
    CompletionOptions,
    CompletionParams,
    DidChangeTextDocumentParams,
    DidOpenTextDocumentParams,
)
from pygls.lsp.types.language_features import (
    CompletionList,
    DeclarationParams,
    DefinitionParams,
    HoverParams,
    List,
    Location,
    Hover
)
from pygls.server import LanguageServer
from pygls.workspace import Document

from src.completions import Completer
from src.navigation import Navigator
from src.utils import extract_enum_name, get_word_at_cursor

from .ast import AST
from .diagnostics import get_diagnostics

server = LanguageServer("vyper", "v0.0.1")
completer = Completer()
navigator = Navigator()
ast = AST()

def validate_doc(ls, params):
    text_doc = ls.workspace.get_document(params.text_document.uri)
    diagnostics = get_diagnostics(text_doc)
    ls.publish_diagnostics(params.text_document.uri, diagnostics)
    ast.update_ast(text_doc)


@server.feature(TEXT_DOCUMENT_DID_OPEN)
async def did_open(ls: LanguageServer, params: DidOpenTextDocumentParams):
    validate_doc(ls, params)


@server.feature(TEXT_DOCUMENT_DID_CHANGE)
async def did_change(ls: LanguageServer, params: DidChangeTextDocumentParams):
    validate_doc(ls, params)


@server.feature(
    TEXT_DOCUMENT_COMPLETION, CompletionOptions(trigger_characters=[":", ".", "@", " "])
)
def completions(ls, params: CompletionParams) -> CompletionList:
    return completer.get_completions(ls, params)


@server.feature(TEXT_DOCUMENT_DECLARATION)
def go_to_declaration(ls: LanguageServer, params: DeclarationParams) -> Optional[Location]:
    document = ls.workspace.get_document(params.text_document.uri)
    range = navigator.find_declaration(document, params.position)
    if range:
        return Location(uri=params.text_document.uri, range=range)

@server.feature(TEXT_DOCUMENT_DEFINITION)
def go_to_definition(ls: LanguageServer, params: DefinitionParams) -> Optional[Location]:
    # TODO: Look for assignment nodes to find definition
    document = ls.workspace.get_document(params.text_document.uri)
    range = navigator.find_declaration(document, params.position)
    if range:
        return Location(uri=params.text_document.uri, range=range)


def get_enum_name(ls: LanguageServer, doc: Document, variant_line_no: int):
    for line_no in range(variant_line_no, 0):
        line = doc.lines[line_no]
        enum_name = extract_enum_name(line)
        if enum_name:
            return enum_name

@server.feature(TEXT_DOCUMENT_REFERENCES)
def find_references(ls: LanguageServer, params: DefinitionParams) -> List[Location]:
    document = ls.workspace.get_document(params.text_document.uri)
    return [
        Location(uri=params.text_document.uri, range=range)
        for range in navigator.find_references(document, params.position)
    ]

@server.feature(HOVER)
def hover(ls: LanguageServer, params: HoverParams):
    document = ls.workspace.get_document(params.text_document.uri)
    hover_info = navigator.get_hover_info(document, params.position)
    if hover_info:
        return Hover(contents=hover_info, range=None)

@server.feature(IMPLEMENTATION)
def implementation(ls: LanguageServer, params: DefinitionParams):
    document = ls.workspace.get_document(params.text_document.uri)
    range = navigator.find_implementation(document, params.position)
    if range:
        return Location(uri=params.text_document.uri, range=range)


def main():
    server.start_io()

src/main.py
vyperlang/vyper-lsp
import re
import sys
from pygls.lsp.types.language_features import Location, Position, Range
from typing import List, Optional

from pygls.workspace import Document
from src.ast import AST
from src.utils import get_expression_at_cursor, get_word_at_cursor

# this class should abstract away all the AST stuff
# and just provide a simple interface for navigation
#
# the navigator should mainly return Ranges
class Navigator:
    def __init__(self, ast=None):
        self.ast = ast or AST()

    def find_state_variable_declaration(self, word: str) -> Optional[Range]:
        node = self.ast.find_state_variable_declaration_node(word)
        if node:
            range = Range(
                start=Position(line=node.lineno - 1, character=node.col_offset),
                end=Position(line=node.end_lineno - 1, character=node.end_col_offset),
            )
            return range

    def find_function_declaration(self, word: str) -> Optional[Range]:
        node = self.ast.find_function_declaration_node_for_name(word)
        if node:
            range = Range(
                start=Position(line=node.lineno - 1, character=node.col_offset),
                end=Position(line=node.end_lineno - 1, character=node.end_col_offset),
            )
            return range

    def find_type_declaration(self, word: str) -> Optional[Range]:
        node = self.ast.find_type_declaration_node(word)
        if node:
            range = Range(
                start=Position(line=node.lineno - 1, character=node.col_offset),
                end=Position(line=node.end_lineno - 1, character=node.end_col_offset),
            )
            return range

    def find_references(self, doc: Document, pos: Position) -> List[Range]:
        print("finding references!!!!", file=sys.stderr)
        og_line = doc.lines[pos.line]
        word = get_word_at_cursor(og_line, pos.character)
        if self.ast.ast_data is None:
            print("ast data is none", file=sys.stderr)
            return []
        references = []

        print("finding references for", word, file=sys.stderr)
        if word in self.ast.get_enums():
            # find all references to this type
            refs = self.ast.find_nodes_referencing_enum(word)
            for ref in refs:
                range = Range(
                    start=Position(line=ref.lineno - 1, character=ref.col_offset),
                    end=Position(line=ref.end_lineno - 1, character=ref.end_col_offset),
                )
                references.append(range)
        elif word in self.ast.get_structs() or word in self.ast.get_events():
            print("found struct or event", word, file=sys.stderr)
            refs = self.ast.find_nodes_referencing_struct(word)
            for ref in refs:
                range = Range(
                    start=Position(line=ref.lineno - 1, character=ref.col_offset),
                    end=Position(line=ref.end_lineno - 1, character=ref.end_col_offset),
                )
                references.append(range)
        elif not og_line[0].isspace() and word in self.ast.get_state_variables():
            print("found state variable", word, file=sys.stderr)
            if "constant(" in og_line:
                refs = self.ast.find_nodes_referencing_constant(word)
            else:
                refs = self.ast.find_nodes_referencing_state_variable(word)
            for ref in refs:
                range = Range(
                    start=Position(line=ref.lineno - 1, character=ref.col_offset),
                    end=Position(line=ref.end_lineno - 1, character=ref.end_col_offset),
                )
                references.append(range)
        return references

    def find_declaration(self, document: Document, pos: Position) -> Optional[Range]:
        if self.ast.ast_data is None:
            return None

        og_line = document.lines[pos.line]
        word = get_word_at_cursor(og_line, pos.character)
        full_word = get_expression_at_cursor(og_line, pos.character)
        range = None

        if full_word.startswith("self."):
            if "(" in full_word:
                range = self.find_function_declaration(word)
            else:
                range = self.find_state_variable_declaration(word)
        else:
            if word in self.ast.get_user_defined_types():
                range = self.find_type_declaration(word)
            elif word in self.ast.get_constants():
                range = self.find_state_variable_declaration(word)
            else:
                # check if full_word matches "enum.variant" regex
                pattern = r"([a-zA-Z_][a-zA-Z0-9_]*)\.([a-zA-Z_][a-zA-Z0-9_]*)"
                match = re.match(pattern, full_word)
                if match:
                    enum_name = match.group(1)
                    variant_name = match.group(2)
                    if enum_name in self.ast.get_enums():
                        variants = self.ast.get_enum_variants(enum_name)
                        if variant_name in variants:
                            range = self.find_type_declaration(enum_name)
        if range:
            return range

    def find_implementation(self, document: Document, pos: Position) -> Optional[Range]:
        og_line = document.lines[pos.line]
        word = get_word_at_cursor(og_line, pos.character)
        expression = get_expression_at_cursor(og_line, pos.character)

        if not "(" in expression:
            return None

        if expression.startswith("self."):
            return self.find_function_declaration(word)
        elif og_line[0].isspace() and og_line.strip().startswith("def"):
            # only lookup external fns if we're in an interface def
            return self.find_function_declaration(word)
        else:
            return None


    def get_hover_info(self, document: Document, pos: Position) -> Optional[str]:
        og_line = document.lines[pos.line]
        word = get_word_at_cursor(og_line, pos.character)
        full_word = get_expression_at_cursor(og_line, pos.character)

        if full_word.startswith("self."):
            if "(" in full_word:
                node = self.ast.find_function_declaration_node_for_name(word)
                if node:
                    fn_name = node.name
                    arg_str = ", ".join([f"{arg.arg}: {arg.annotation.id}" for arg in node.args.args])
                    return f"(Internal Function) **{fn_name}**({arg_str})"
            else:
                node = self.ast.find_state_variable_declaration_node(word)
                if node:
                    variable_type = node.annotation.id
                    return f"(State Variable) **{word}** : **{variable_type}**"
        else:
            if word in self.ast.get_structs():
                node = self.ast.find_type_declaration_node(word)
                if node:
                    return f"(Struct) **{word}**"
            elif word in self.ast.get_enums():
                node = self.ast.find_type_declaration_node(word)
                if node:
                    return f"(Enum) **{word}**"
            elif word in self.ast.get_events():
                node = self.ast.find_type_declaration_node(word)
                if node:
                    return f"(Event) **{word}**"
            elif word in self.ast.get_constants():
                node = self.ast.find_state_variable_declaration_node(word)
                if node:
                    variable_type = node.annotation.id
                    return f"(Constant) **{word}** : **{variable_type}**"
            else:
                return None

src/navigation.py
vyperlang/vyper-lsp
from pathlib import Path

import vyper
from lark import Lark
from vyper.ast.grammar import PythonIndenter

GRAMMAR_FILE_PATH = Path(vyper.__file__).parent / "ast" / "grammar.lark"
GRAMMAR = GRAMMAR_FILE_PATH.read_text()

parser = Lark(GRAMMAR, parser="lalr", start="module", postlex=PythonIndenter())

src/parse.py
vyperlang/vyper-lsp
# Available base types
UNSIGNED_INTEGER_TYPES = {f"uint{8*(i+1)}" for i in range(32)}
SIGNED_INTEGER_TYPES = {f"int{8*(i+1)}" for i in range(32)}
INTEGER_TYPES = UNSIGNED_INTEGER_TYPES | SIGNED_INTEGER_TYPES

BYTES_M_TYPES = {f"bytes{i+1}" for i in range(32)}
DECIMAL_TYPES = {"decimal"}


BASE_TYPES = INTEGER_TYPES | BYTES_M_TYPES | DECIMAL_TYPES | {"bool", "address"}

src/types.py
vyperlang/vyper-lsp
import re
from pathlib import Path

from vyper.compiler import CompilerData


def get_source(filepath):
    base_path = Path(__file__).parent.parent
    filepath = base_path / filepath
    return filepath.read_text()


def get_compiler_data(filepath):
    source = get_source(filepath)
    return CompilerData(source)


# detect if current line is a variable declaration
def is_var_declaration(line):
    # regex for variable declaration
    # should match lines starting with any identifier followed by a colon
    # like "foo: "
    reg = r"^\s*[a-zA-Z_][a-zA-Z0-9_]*\s*:"
    return bool(re.match(reg, line.strip()))


def is_attribute_access(line):
    # regex for attribute access
    # should match lines ending with a dot
    # like "foo."
    reg = r"\s*\.\s*$"
    return bool(re.match(reg, line.strip()))


def get_word_at_cursor(sentence: str, cursor_index: int) -> str:
    start = cursor_index
    end = cursor_index

    # Find the start of the word
    while start > 0 and sentence[start - 1].isalnum():
        start -= 1

    # Find the end of the word
    while end < len(sentence) and sentence[end].isalnum():
        end += 1

    # Extract the word
    word = sentence[start:end]

    return word

def get_expression_at_cursor(sentence: str, cursor_index: int) -> str:
    # does the same thing as get_word_at_cursor but includes . and [ and ] in the expression
    start = cursor_index
    end = cursor_index

    # Find the start of the word
    while start > 0 and sentence[start - 1].isalnum() or sentence[start - 1] in ".[]()":
        start -= 1

    # Find the end of the word
    while end < len(sentence) and sentence[end].isalnum() or sentence[end] in ".[]()":
        end += 1

    # Extract the word
    word = sentence[start:end]

    return word

def extract_enum_name(line: str):
    match = re.match(r"enum\s+([a-zA-Z_][a-zA-Z0-9_]*)\s*:", line)
    if match:
        return match.group(1)
    return None

src/utils.py
vyperlang/vyper-lsp