address comments; remove LSPBaseReference class

Author: themisvaltinos

sqlmesh/lsp/main.py

@@ -44,7 +44,12 @@
     CustomMethod,
 )
 from sqlmesh.lsp.hints import get_hints
-from sqlmesh.lsp.reference import LSPModelReference, get_references, get_all_references
+from sqlmesh.lsp.reference import (
+    LSPCteReference,
+    LSPModelReference,
+    get_references,
+    get_all_references,
+)
 from sqlmesh.lsp.uri import URI
 from web.server.api.endpoints.lineage import column_lineage, model_lineage
 from web.server.api.endpoints.models import get_models
@@ -368,7 +373,7 @@ def hover(ls: LanguageServer, params: types.HoverParams) -> t.Optional[types.Hov
                 if not references:
                     return None
                 reference = references[0]
-                if not reference.markdown_description:
+                if isinstance(reference, LSPCteReference) or not reference.markdown_description:
                     return None
                 return types.Hover(
                     contents=types.MarkupContent(
@@ -419,7 +424,7 @@ def goto_definition(
                 location_links = []
                 for reference in references:
                     # Use target_range if available (CTEs, Macros), otherwise default to start of file
-                    if not isinstance(reference, LSPModelReference) and reference.target_range:
+                    if not isinstance(reference, LSPModelReference):
                         target_range = reference.target_range
                         target_selection_range = reference.target_range
                     else:

sqlmesh/lsp/reference.py

@@ -24,32 +24,32 @@
 import inspect
 
 
-class LSPBaseReference(PydanticModel):
-    """Base class for all LSP reference types."""
-
-    range: Range
-    uri: str
-    markdown_description: t.Optional[str] = None
-
-
-class LSPModelReference(LSPBaseReference):
+class LSPModelReference(PydanticModel):
     """A LSP reference to a model."""
 
     type: t.Literal["model"] = "model"
+    uri: str
+    range: Range
+    markdown_description: t.Optional[str] = None
 
 
-class LSPCteReference(LSPBaseReference):
+class LSPCteReference(PydanticModel):
     """A LSP reference to a CTE."""
 
     type: t.Literal["cte"] = "cte"
+    uri: str
+    range: Range
     target_range: Range
 
 
-class LSPMacroReference(LSPBaseReference):
+class LSPMacroReference(PydanticModel):
     """A LSP reference to a macro."""
 
     type: t.Literal["macro"] = "macro"
+    uri: str
+    range: Range
     target_range: Range
+    markdown_description: t.Optional[str] = None
 
 
 Reference = t.Annotated[
@@ -451,12 +451,16 @@ def get_model_find_all_references(
         A list of references to the model across all files
     """
     # First, get the references in the current file to determine what model we're looking for
-    current_file_references = get_model_definitions_for_a_path(lint_context, document_uri)
+    current_file_references = [
+        ref
+        for ref in get_model_definitions_for_a_path(lint_context, document_uri)
+        if isinstance(ref, LSPModelReference)
+    ]
 
     # Find the model reference at the cursor position
     target_model_uri: t.Optional[str] = None
     for ref in current_file_references:
-        if _position_within_range(position, ref.range) and isinstance(ref, LSPModelReference):
+        if _position_within_range(position, ref.range):
             # This is a model reference, get the target model URI
             target_model_uri = ref.uri
             break
@@ -499,7 +503,11 @@ def get_model_find_all_references(
             continue
 
         # Get model references for this file
-        file_references = get_model_definitions_for_a_path(lint_context, file_uri)
+        file_references = [
+            ref
+            for ref in get_model_definitions_for_a_path(lint_context, file_uri)
+            if isinstance(ref, LSPModelReference)
+        ]
 
         # Add references that point to the target model file
         for ref in file_references:
@@ -562,7 +570,6 @@ def get_cte_references(
             uri=document_uri.value,
             range=target_cte_definition_range,
             target_range=target_cte_definition_range,
-            markdown_description="CTE definition",
         )
     ]
 
@@ -574,7 +581,6 @@ def get_cte_references(
                     uri=document_uri.value,
                     range=ref.range,
                     target_range=ref.target_range,
-                    markdown_description="CTE usage",
                 )
             )
 

tests/lsp/test_reference.py

@@ -1,7 +1,7 @@
 from lsprotocol.types import Position
 from sqlmesh.core.context import Context
 from sqlmesh.lsp.context import LSPContext, ModelTarget, AuditTarget
-from sqlmesh.lsp.reference import get_model_definitions_for_a_path, by_position
+from sqlmesh.lsp.reference import LSPModelReference, get_model_definitions_for_a_path, by_position
 from sqlmesh.lsp.uri import URI
 
 
@@ -47,9 +47,13 @@ def test_reference_with_alias() -> None:
         if isinstance(info, ModelTarget) and "sushi.waiter_revenue_by_day" in info.names
     )
 
-    references = get_model_definitions_for_a_path(
-        lsp_context, URI.from_path(waiter_revenue_by_day_path)
-    )
+    references = [
+        ref
+        for ref in get_model_definitions_for_a_path(
+            lsp_context, URI.from_path(waiter_revenue_by_day_path)
+        )
+        if isinstance(ref, LSPModelReference)
+    ]
     assert len(references) == 3
 
     with open(waiter_revenue_by_day_path, "r") as file:

tests/lsp/test_reference_cte.py

@@ -1,7 +1,7 @@
 import re
 from sqlmesh.core.context import Context
 from sqlmesh.lsp.context import LSPContext, ModelTarget
-from sqlmesh.lsp.reference import get_references
+from sqlmesh.lsp.reference import LSPCteReference, get_references
 from sqlmesh.lsp.uri import URI
 from lsprotocol.types import Range, Position
 import typing as t
@@ -28,7 +28,7 @@ def test_cte_parsing():
     references = get_references(lsp_context, URI.from_path(sushi_customers_path), position)
     assert len(references) == 1
     assert references[0].uri == URI.from_path(sushi_customers_path).value
-    assert references[0].markdown_description is None
+    assert isinstance(references[0], LSPCteReference)
     assert (
         references[0].range.start.line == ranges[1].start.line
     )  # The reference location (where we clicked)
@@ -43,7 +43,7 @@ def test_cte_parsing():
     references = get_references(lsp_context, URI.from_path(sushi_customers_path), position)
     assert len(references) == 1
     assert references[0].uri == URI.from_path(sushi_customers_path).value
-    assert references[0].markdown_description is None
+    assert isinstance(references[0], LSPCteReference)
     assert (
         references[0].range.start.line == ranges[1].start.line
     )  # The reference location (where we clicked)