feat: add inlay type hints for sqlmesh sql models

Author: Christopher Giroir

sqlmesh/lsp/hints.py

@@ -0,0 +1,113 @@
+"""Type hinting on SQLMesh models"""
+
+import typing as t
+
+from lsprotocol import types
+
+from sqlglot import exp
+from sqlglot.expressions import Expression
+from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
+from sqlglot.optimizer.scope import build_scope, find_all_in_scope
+from sqlmesh.core.model.definition import SqlModel
+from sqlmesh.lsp.context import LSPContext, ModelTarget
+from sqlmesh.lsp.uri import URI
+
+
+def get_hints(
+    lsp_context: LSPContext,
+    document_uri: URI,
+    start_line: int,
+    end_line: int,
+) -> t.List[types.InlayHint]:
+    """
+    Get type hints for certain lines in a document
+
+    Args:
+        lint_context: The LSP context
+        document_uri: The URI of the document
+        start_line: the starting line to get hints for
+        end_line: the ending line to get hints for
+
+    Returns:
+        A list of hints to apply to the document
+    """
+    path = document_uri.to_path()
+    if path.suffix != ".sql":
+        return []
+
+    if path not in lsp_context.map:
+        return []
+
+    file_info = lsp_context.map[path]
+
+    # Process based on whether it's a model or standalone audit
+    if isinstance(file_info, ModelTarget):
+        # It's a model
+        model = lsp_context.context.get_model(
+            model_or_snapshot=file_info.names[0], raise_if_missing=False
+        )
+        if model is None or not isinstance(model, SqlModel):
+            return []
+
+        query = model.query
+        dialect = model.dialect
+        columns_to_types = model.columns_to_types or {}
+    else:
+        return []
+
+    return _get_type_hints_for_model_from_query(
+        query, dialect, columns_to_types, start_line, end_line
+    )
+
+
+def _get_type_hints_for_model_from_query(
+    query: Expression,
+    dialect: str,
+    columns_to_types: t.Dict[str, t.Any],
+    start_line: int,
+    end_line: int,
+) -> t.List[types.InlayHint]:
+    hints: t.List[types.InlayHint] = []
+    try:
+        query = normalize_identifiers(query.copy(), dialect=dialect)
+        root = build_scope(query)
+
+        if not root:
+            return []
+
+        for select in find_all_in_scope(root.expression, exp.Select):
+            for select_exp in select.expressions:
+                if not select_exp:
+                    continue
+
+                if isinstance(select_exp, exp.Alias):
+                    meta = select_exp.args["alias"]._meta
+                elif isinstance(select_exp, exp.Column):
+                    meta = select_exp.parts[-1]._meta
+                else:
+                    continue
+
+                line = meta.get("line") - 1  # Lines from sqlglot are 1 based
+                col = meta.get("col")
+
+                name = select_exp.alias_or_name
+                if name not in columns_to_types:
+                    continue
+
+                if line < start_line or line > end_line:
+                    continue
+
+                type_label = str(columns_to_types.get(name))
+                hints.append(
+                    types.InlayHint(
+                        label=f"::{type_label}",
+                        kind=types.InlayHintKind.Type,
+                        padding_left=False,
+                        padding_right=True,
+                        position=types.Position(line=line, character=col),
+                    )
+                )
+
+        return hints
+    except Exception:
+        return []

sqlmesh/lsp/main.py

@@ -36,9 +36,8 @@
     RenderModelRequest,
     RenderModelResponse,
 )
-from sqlmesh.lsp.reference import (
-    get_references,
-)
+from sqlmesh.lsp.hints import get_hints
+from sqlmesh.lsp.reference import get_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
@@ -324,6 +323,27 @@ def hover(ls: LanguageServer, params: types.HoverParams) -> t.Optional[types.Hov
                 ls.show_message(f"Error getting hover information: {e}", types.MessageType.Error)
                 return None
 
+        @self.server.feature(types.TEXT_DOCUMENT_INLAY_HINT)
+        def inlay_hint(
+            ls: LanguageServer, params: types.InlayHintParams
+        ) -> t.List[types.InlayHint]:
+            """Implement type hints for sql columns as inlay hints"""
+            try:
+                uri = URI(params.text_document.uri)
+                self._ensure_context_for_document(uri)
+                document = ls.workspace.get_text_document(params.text_document.uri)
+                if self.lsp_context is None:
+                    raise RuntimeError(f"No context found for document: {document.path}")
+
+                start_line = params.range.start.line
+                end_line = params.range.end.line
+                hints = get_hints(self.lsp_context, uri, start_line, end_line)
+                return hints
+
+            except Exception as e:
+                ls.show_message(f"Error getting type hints: {e}", types.MessageType.Error)
+                return []
+
         @self.server.feature(types.TEXT_DOCUMENT_DEFINITION)
         def goto_definition(
             ls: LanguageServer, params: types.DefinitionParams

tests/lsp/test_hints.py

@@ -0,0 +1,62 @@
+"""Tests for type hinting SQLMesh models"""
+
+import pytest
+
+from sqlmesh.core.context import Context
+from sqlmesh.lsp.context import LSPContext, ModelTarget
+from sqlmesh.lsp.hints import get_hints
+from sqlmesh.lsp.uri import URI
+
+
+@pytest.mark.fast
+def test_hints() -> None:
+    context = Context(paths=["examples/sushi"])
+    lsp_context = LSPContext(context)
+
+    # Find model URIs
+    active_customers_path = next(
+        path
+        for path, info in lsp_context.map.items()
+        if isinstance(info, ModelTarget) and "sushi.active_customers" in info.names
+    )
+    customer_revenue_lifetime_path = next(
+        path
+        for path, info in lsp_context.map.items()
+        if isinstance(info, ModelTarget) and "sushi.customer_revenue_lifetime" in info.names
+    )
+    customer_revenue_by_day_path = next(
+        path
+        for path, info in lsp_context.map.items()
+        if isinstance(info, ModelTarget) and "sushi.customer_revenue_by_day" in info.names
+    )
+
+    active_customers_uri = URI.from_path(active_customers_path)
+    ac_hints = get_hints(lsp_context, active_customers_uri, start_line=0, end_line=9999)
+    assert len(ac_hints) == 2
+    assert ac_hints[0].label == "::INT"
+    assert ac_hints[1].label == "::TEXT"
+
+    customer_revenue_lifetime_uri = URI.from_path(customer_revenue_lifetime_path)
+    crl_hints = get_hints(
+        lsp_context=lsp_context,
+        document_uri=customer_revenue_lifetime_uri,
+        start_line=0,
+        end_line=9999,
+    )
+    assert len(crl_hints) == 3
+    assert crl_hints[0].label == "::INT"
+    assert crl_hints[1].label == "::DOUBLE"
+    assert crl_hints[2].label == "::DATE"
+
+    customer_revenue_by_day_uri = URI.from_path(customer_revenue_by_day_path)
+    crbd_hints = get_hints(
+        lsp_context=lsp_context,
+        document_uri=customer_revenue_by_day_uri,
+        start_line=0,
+        end_line=9999,
+    )
+    assert len(crbd_hints) == 4
+    assert crbd_hints[0].label == "::INT"
+    assert crbd_hints[1].label == "::DOUBLE"
+    assert crbd_hints[2].label == "::INT"
+    assert crbd_hints[3].label == "::DATE"