feat(lsp): include macro details in completions

Author: benfdking

sqlmesh/lsp/completions.py

@@ -1,6 +1,6 @@
 from functools import lru_cache
 from sqlglot import Dialect, Tokenizer
-from sqlmesh.lsp.custom import AllModelsResponse
+from sqlmesh.lsp.custom import AllModelsResponse, MacroCompletion
 from sqlmesh import macro
 import typing as t
 from sqlmesh.lsp.context import AuditTarget, LSPContext, ModelTarget
@@ -60,15 +60,23 @@ def get_models(context: t.Optional[LSPContext], file_uri: t.Optional[URI]) -> t.
     return all_models
 
 
-def get_macros(context: t.Optional[LSPContext], file_uri: t.Optional[URI]) -> t.Set[str]:
-    """Return a set of all macros with the ``@`` prefix."""
-    names = set(macro.get_registry())
+def get_macros(
+    context: t.Optional[LSPContext], file_uri: t.Optional[URI]
+) -> t.List[MacroCompletion]:
+    """Return a list of macros with optional descriptions."""
+    macros: t.Dict[str, t.Optional[str]] = {}
+
+    for name, m in macro.get_registry().items():
+        macros[name] = getattr(m.func, "__doc__", None)
+
     try:
         if context is not None:
-            names.update(context.context._macros)
+            for name, m in context.context._macros.items():
+                macros[name] = getattr(m.func, "__doc__", None)
     except Exception:
         pass
-    return names
+
+    return [MacroCompletion(name=name, description=doc) for name, doc in macros.items()]
 
 
 def get_keywords(context: t.Optional[LSPContext], file_uri: t.Optional[URI]) -> t.Set[str]:

sqlmesh/lsp/custom.py

@@ -23,14 +23,21 @@ class AllModelsRequest(CustomMethodRequestBaseClass):
     textDocument: types.TextDocumentIdentifier
 
 
+class MacroCompletion(PydanticModel):
+    """Information about a macro for autocompletion."""
+
+    name: str
+    description: t.Optional[str] = None
+
+
 class AllModelsResponse(CustomMethodResponseBaseClass):
     """
     Response to get all the models that are in the current project.
     """
 
     models: t.List[str]
     keywords: t.List[str]
-    macros: t.List[str]
+    macros: t.List["MacroCompletion"]
 
 
 RENDER_MODEL_FEATURE = "sqlmesh/render_model"

sqlmesh/lsp/main.py

@@ -625,7 +625,8 @@ def completion(
                     and getattr(params.context, "trigger_character", None) == "@"
                 )
 
-                for macro_name in completion_response.macros:
+                for macro in completion_response.macros:
+                    macro_name = macro.name
                     insert_text = macro_name if triggered_by_at else f"@{macro_name}"
 
                     completion_items.append(
@@ -636,6 +637,7 @@ def completion(
                             filter_text=macro_name,
                             kind=types.CompletionItemKind.Function,
                             detail="SQLMesh Macro",
+                            documentation=macro.description,
                         )
                     )
 

tests/lsp/test_completions.py

@@ -33,8 +33,9 @@ def test_get_macros():
     file_uri = URI.from_path(file_path)
     completions = LSPContext.get_completions(lsp_context, file_uri, file_content)
 
-    assert "each" in completions.macros
-    assert "add_one" in completions.macros
+    macro_names = {m.name for m in completions.macros}
+    assert "each" in macro_names
+    assert "add_one" in macro_names
 
 
 def test_get_sql_completions_with_context_no_file_uri():