feat: improve suggestion of internal fns & state

tests/conftest.py

@@ -0,0 +1,8 @@
+import pytest
+
+from vyper_lsp.ast import AST
+
+
+@pytest.fixture
+def ast():
+    return AST()

tests/test_ast.py

@@ -1,14 +1,7 @@
 from pygls.lsp.types import Position
-import pytest
-
 from vyper_lsp.ast import AST
 
 
-@pytest.fixture
-def ast():
-    return AST()
-
-
 def test_get_constants(ast):
     src = """
 x: constant(uint256) = 123

tests/test_completions.py

@@ -0,0 +1,35 @@
+from pygls.lsp.types import CompletionContext, CompletionParams, Position
+from pygls.workspace import Document
+
+from vyper_lsp.analyzer.AstAnalyzer import AstAnalyzer
+
+
+def test_completion(ast):
+    src = """
+@internal
+def foo():
+    return
+
+@external
+def bar():
+    self.foo()
+"""
+    ast.build_ast(src)
+
+    src += """
+@external
+def baz():
+    self.
+"""
+
+    doc = Document(uri="examples/Foo.vy", source=src)
+    pos = Position(line=11, character=7)
+    context = CompletionContext(trigger_character=".", trigger_kind=2)
+    params = CompletionParams(
+        text_document={"uri": doc.uri, "source": src}, position=pos, context=context
+    )
+
+    analyzer = AstAnalyzer(ast)
+    completions = analyzer.get_completions_in_doc(doc, params)
+    assert len(completions.items) == 1
+    assert "foo" in [c.label for c in completions.items]

vyper_lsp/analyzer/AstAnalyzer.py

@@ -7,6 +7,7 @@
 from vyper.compiler import CompilerData
 from vyper.exceptions import VyperException
 from vyper_lsp.analyzer.BaseAnalyzer import Analyzer
+from vyper_lsp.ast import AST
 from vyper_lsp.utils import (
     get_expression_at_cursor,
     get_word_at_cursor,
@@ -38,28 +39,33 @@
 
 
 class AstAnalyzer(Analyzer):
-    def __init__(self, ast) -> None:
+    def __init__(self, ast: AST) -> None:
         super().__init__()
         self.ast = ast
         if get_installed_vyper_version() < min_vyper_version:
             self.diagnostics_enabled = False
         else:
             self.diagnostics_enabled = True
 
-    def get_completions(
-        self, ls: LanguageServer, params: CompletionParams
+    def get_completions_in_doc(
+        self, document: Document, params: CompletionParams
     ) -> CompletionList:
         items = []
-        document = ls.workspace.get_document(params.text_document.uri)
         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))
+                if element == "self":
+                    for fn in self.ast.get_internal_functions():
+                        items.append(CompletionItem(label=fn))
+                    for var in self.ast.get_state_variables():
+                        items.append(CompletionItem(label=var))
+                else:
+                    for attr in self.ast.get_attributes_for_symbol(element):
+                        items.append(CompletionItem(label=attr))
                 completions = CompletionList(is_incomplete=False, items=[])
                 completions.add_items(items)
                 return completions
@@ -90,6 +96,12 @@ def get_completions(
         else:
             return CompletionList(is_incomplete=False, items=[])
 
+    def get_completions(
+        self, ls: LanguageServer, params: CompletionParams
+    ) -> CompletionList:
+        document = ls.workspace.get_document(params.text_document.uri)
+        return self.get_completions_in_doc(document, params)
+
     def hover_info(self, document: Document, pos: Position) -> Optional[str]:
         if len(document.lines) < pos.line:
             return None