Moving more functions to pydantic

Author: yeger00

pylspclient/lsp_client.py

@@ -1,13 +1,17 @@
+from typing import Optional
+
 from pydantic import ValidationError
 from pylspclient import lsp_structs
-from pylspclient.lsp_pydantic_strcuts import TextDocumentItem, TextDocumentIdentifier, DocumentSymbol, SymbolInformation
+from pylspclient.lsp_endpoint import LspEndpoint
+from pylspclient.lsp_pydantic_strcuts import TextDocumentItem, TextDocumentIdentifier, DocumentSymbol, SymbolInformation, LocationLink, Location
+from pylspclient.lsp_pydantic_strcuts import Position
 
 class LspClient(object):
-    def __init__(self, lsp_endpoint):
+    def __init__(self, lsp_endpoint: LspEndpoint):
         """
         Constructs a new LspClient instance.
 
-        :param lsp_endpoint: TODO
+        :param lsp_endpoint:
         """
         self.lsp_endpoint = lsp_endpoint
 
@@ -85,7 +89,7 @@ def didOpen(self, textDocument: TextDocumentItem):
 
         :param TextDocumentItem textDocument: The document that was opened.
         """
-        return self.lsp_endpoint.send_notification("textDocument/didOpen", textDocument=textDocument.dict())
+        return self.lsp_endpoint.send_notification("textDocument/didOpen", textDocument=textDocument)
 
 
     def didChange(self, textDocument, contentChanges):
@@ -98,7 +102,7 @@ def didChange(self, textDocument, contentChanges):
             to the document. So if there are two content changes c1 and c2 for a document in state S then c1 move the document 
             to S' and c2 to S''.
         """
-        return self.lsp_endpoint.send_notification("textDocument/didChange", textDocument=textDocument, contentChanges=contentChanges)
+        self.lsp_endpoint.send_notification("textDocument/didChange", textDocument=textDocument, contentChanges=contentChanges)
 
 
     def documentSymbol(self, textDocument: TextDocumentIdentifier) -> list[DocumentSymbol] | list[SymbolInformation]:
@@ -115,15 +119,19 @@ def documentSymbol(self, textDocument: TextDocumentIdentifier) -> list[DocumentS
             return [SymbolInformation.parse_obj(sym) for sym in result_dict]
 
 
-    def typeDefinition(self, textDocument, position):
+    def typeDefinition(
+        self,
+        textDocument: TextDocumentIdentifier,
+        position: Position
+    ) -> list[Location]:
         """
         The goto type definition request is sent from the client to the server to resolve the type definition location of a symbol at a given text document position.
 
         :param TextDocumentItem textDocument: The text document.
         :param Position position: The position inside the text document.
         """
         result_dict = self.lsp_endpoint.call_method("textDocument/typeDefinition", textDocument=textDocument, position=position)
-        return [lsp_structs.Location(**result) for result in result_dict]
+        return [Location.parse_obj(result) for result in result_dict]
 
 
     def signatureHelp(self, textDocument, position):
@@ -153,7 +161,11 @@ def completion(self, textDocument, position, context):
             return [lsp_structs.CompletionItem(**result) for result in result_dict]
 
 
-    def declaration(self, textDocument, position):
+    def declaration(
+        self,
+        textDocument: TextDocumentIdentifier,
+        position: Position
+    ) -> Location | list[Location] | list[LocationLink]:
             """
             The go to declaration request is sent from the client to the server to resolve the declaration location of a 
             symbol at a given text document position.
@@ -166,12 +178,18 @@ def declaration(self, textDocument, position):
             """
             result_dict = self.lsp_endpoint.call_method("textDocument/declaration", textDocument=textDocument, position=position)
             if "uri" in result_dict:
-                return lsp_structs.Location(**result_dict)
+                return Location.parse_obj(result_dict)
 
-            return [lsp_structs.Location(**result) if "uri" in result else lsp_structs.LocationLink(**result) for result in result_dict]
+            return [Location.parse_obj(result) if "uri" in result else LocationLink.parse_obj(result) for result in result_dict]
 
 
-    def definition(self, textDocument, position):
+    def definition(
+        self,
+        textDocument: TextDocumentIdentifier,
+        position: Position,
+        workDoneToken: Optional[str] = None,
+        partialResultToken: Optional[str] = None
+    ) -> Location | list[Location] | list[LocationLink]:
             """
             The go to definition request is sent from the client to the server to resolve the declaration location of a 
             symbol at a given text document position.
@@ -184,6 +202,6 @@ def definition(self, textDocument, position):
             """
             result_dict = self.lsp_endpoint.call_method("textDocument/definition", textDocument=textDocument, position=position)
             if "uri" in result_dict:
-                return lsp_structs.Location(**result_dict)
+                return Location.parse_obj(result_dict)
 
-            return [lsp_structs.Location(**result) if "uri" in result else lsp_structs.LocationLink(**result) for result in result_dict]
+            return [Location.parse_obj(result) if "uri" in result else LocationLink.parse_obj(result) for result in result_dict]

pylspclient/lsp_pydantic_strcuts.py

@@ -1,6 +1,6 @@
 from typing import Optional, List
 from enum import Enum, IntEnum
-from pydantic import BaseModel
+from pydantic import BaseModel, HttpUrl
 
 
 class LanguageIdentifier(str, Enum):
@@ -146,3 +146,30 @@ class SymbolInformation(BaseModel):
     deprecated: Optional[bool] = None
     location: Location
     containerName: Optional[str] = None
+
+class TextDocumentPositionParams(BaseModel):
+    """
+    A base class including the text document identifier and a position within that document.
+    """
+    textDocument: TextDocumentIdentifier
+    position: Position
+
+class ReferenceContext(BaseModel):
+    """
+    Additional information about the context of a reference request.
+    """
+    includeDeclaration: bool  # Whether to include the declaration of the symbol being referenced
+
+class ReferenceParams(TextDocumentPositionParams):
+    """
+    Parameters for a Reference Request in the Language Server Protocol.
+    """
+    context: ReferenceContext
+    workDoneToken: Optional[str] = None  # Optional; used for progress reporting
+    partialResultToken: Optional[str] = None  # Optional; used for partial results
+
+class LocationLink(BaseModel):
+    originSelectionRange: Optional[Range]
+    targetUri: HttpUrl
+    targetRange: Range
+    targetSelectionRange: Range

tests/test-workspace/lsp_client.py

@@ -1,7 +1,8 @@
 from pylspclient import lsp_structs
+from pylspclient.lsp_endpoint import LspEndpoint
 
 class LspClient(object):
-    def __init__(self, lsp_endpoint):
+    def __init__(self, lsp_endpoint: LspEndpoint):
         """
         Constructs a new LspClient instance.
 

tests/test_pylsp_integration.py

@@ -1,16 +1,23 @@
-import os.path
+from typing import Optional
+from os import path, listdir
 import pytest
 import subprocess
 import threading
 
 import pylspclient
-from pylspclient.lsp_pydantic_strcuts import TextDocumentIdentifier, TextDocumentItem, LanguageIdentifier
+from pylspclient.lsp_pydantic_strcuts import TextDocumentIdentifier, TextDocumentItem, LanguageIdentifier, Position, Range
 
 
 def to_uri(path: str) -> str:
+    if path.startswith("uri://"):
+        return path
     return f"uri://{path}"
 
 
+def from_uri(path: str) -> str:
+    return path.replace("uri://", "").replace("uri:", "")
+
+
 class ReadPipe(threading.Thread):
     def __init__(self, pipe):
         threading.Thread.__init__(self)
@@ -43,7 +50,7 @@ def server_process() -> subprocess.Popen:
         }
     }
 }
-DEFAULT_ROOT = "./tests/test-workspace/"
+DEFAULT_ROOT = path.abspath("./tests/test-workspace/")
 
 
 @pytest.fixture
@@ -89,9 +96,9 @@ def test_initialize(json_rpc: pylspclient.JsonRpcEndpoint):
     lsp_client.exit()
 
 
-def test_type_definition(lsp_client: pylspclient.LspClient):
+def test_definition(lsp_client: pylspclient.LspClient):
     file_path = "lsp_client.py"
-    relative_file_path = os.path.join(DEFAULT_ROOT, file_path)
+    relative_file_path = path.join(DEFAULT_ROOT, file_path)
     uri = to_uri(relative_file_path)
     text = open(relative_file_path, "r").read()
     languageId = LanguageIdentifier.PYTHON
@@ -123,3 +130,49 @@ def test_type_definition(lsp_client: pylspclient.LspClient):
         'documentSymbol'
     ]
     assert set(symbol.name for symbol in symbols) == set(expected_symbols)
+
+
+def add_dir(lsp_client: pylspclient.LspClient, root: str):
+    for filename in listdir(root):
+        if filename.endswith(".py"):
+            add_file(lsp_client, path.join(root, filename))
+
+
+def add_file(lsp_client: pylspclient.LspClient, relative_file_path: str):
+    uri = to_uri(relative_file_path)
+    text = open(relative_file_path, "r").read()
+    languageId = LanguageIdentifier.PYTHON
+    version = 1
+    # First need to open the file, and then iterate over the docuemnt's symbols
+    lsp_client.didOpen(TextDocumentItem(uri=uri, languageId=languageId, version=version, text=text))
+
+
+def string_in_text_to_position(text: str, string: str) -> Optional[Position]:
+    for i, line in enumerate(text.splitlines()):
+        char = line.find(string)
+        if char != -1:
+            return Position(line=i, character=char)
+    return None
+
+
+def range_in_text_to_string(text: str, range_: Range) -> Optional[str]:
+    lines = text.splitlines()
+    if range_.start.line == range_.end.line:
+        # Same line
+        return lines[range_.start.line][range_.start.character:range_.end.character]
+    raise NotImplementedError
+
+
+def test_type_definition(lsp_client: pylspclient.LspClient):
+    add_dir(lsp_client, DEFAULT_ROOT)
+    file_path = "lsp_client.py"
+    relative_file_path = path.join(DEFAULT_ROOT, file_path)
+    uri = to_uri(relative_file_path)
+    file_content = open(relative_file_path, "r").read()
+    position = string_in_text_to_position(file_content, "send_notification")
+    definitions = lsp_client.definition(TextDocumentIdentifier(uri=uri), position)
+    assert len(definitions) == 1
+    result_path = from_uri(definitions[0].uri)
+    result_file_content = open(result_path, "r").read()
+    result_definition = range_in_text_to_string(result_file_content, definitions[0].range)
+    assert result_definition == "send_notification"