Support inlay hints in code chunks

apps/lsp/src/middleware.ts

@@ -28,7 +28,8 @@ export function middlewareCapabilities() : ServerCapabilities {
     },
     documentFormattingProvider: true,
     documentRangeFormattingProvider: true,
-    definitionProvider: true
+    definitionProvider: true,
+    inlayHintProvider: true
   }
 };
 

apps/vscode/src/lsp/client.ts

@@ -62,6 +62,7 @@ import {
 } from "../providers/format";
 import { getHover, getSignatureHelpHover } from "../core/hover";
 import { imageHover } from "../providers/hover-image";
+import { embeddedInlayHintsProvider } from "../providers/inlay-hints";
 import { LspInitializationOptions, QuartoContext } from "quarto-core";
 import { extensionHost } from "../host";
 import semver from "semver";
@@ -104,6 +105,7 @@ export async function activateLsp(
     provideDocumentRangeFormattingEdits: embeddedDocumentRangeFormattingProvider(
       engine
     ),
+    provideInlayHints: embeddedInlayHintsProvider(engine),
   };
   if (config.get("cells.hoverHelp.enabled", true)) {
     middleware.provideHover = embeddedHoverProvider(engine);

apps/vscode/src/providers/inlay-hints.ts

@@ -0,0 +1,99 @@
+import {
+  InlayHint,
+  Range,
+  CancellationToken,
+  commands,
+  TextDocument,
+  Uri
+} from "vscode";
+import { ProvideInlayHintsSignature } from "vscode-languageclient";
+import { unadjustedPosition, adjustedRange, languages, virtualDocForLanguage, withVirtualDocUri, unadjustedRange } from "../vdoc/vdoc";
+import { MarkdownEngine } from "../markdown/engine";
+import { isQuartoDoc, isQuartoYaml } from "../core/doc";
+
+/**
+ * Provides inlay hints for all embedded languages within a single document
+ *
+ * Note that `vscode.executeInlayHintProvider` does not currently "resolve"
+ * inlay hints, so if the underlying provider delays the text edits (that you'd
+ * get when you double click) to resolve time, then we will never see them in the
+ * quarto document (for example, pylance delays, but basedpyright does not).
+ * https://github.com/microsoft/vscode/issues/249359
+ */
+export function embeddedInlayHintsProvider(engine: MarkdownEngine) {
+  return async (
+    document: TextDocument,
+    range: Range,
+    token: CancellationToken,
+    next: ProvideInlayHintsSignature
+  ): Promise<InlayHint[] | null | undefined> => {
+    if (isQuartoYaml(document)) {
+      // The LSP client tracks quarto related yaml files like `_quarto.yaml`,
+      // but we don't provide inlay hints for these. Calling `next()` results
+      // in an "unhandled method" toast notification, so we return `undefined`
+      // directly instead. Is there a better solution?
+      return undefined;
+    }
+    if (!isQuartoDoc(document, true)) {
+      return await next(document, range, token);
+    }
+
+    const tokens = engine.parse(document);
+
+    // Determine all embedded languages used within this document
+    const embeddedLanguages = languages(tokens);
+
+    // Fetch inlay hints for each embedded language
+    const hints: InlayHint[] = [];
+
+    for (const embeddedLanguage of embeddedLanguages) {
+      const vdoc = virtualDocForLanguage(document, tokens, embeddedLanguage);
+
+      if (!vdoc) {
+        const language = embeddedLanguage.ids.at(0) ?? "??";
+        console.error(`[InlayHints] No virtual document produced for language: ${language}.`);
+        continue;
+      };
+
+      // Map `range` into adjusted vdoc range
+      const vdocRange = adjustedRange(vdoc.language, range);
+
+      // Get inlay hints for this embedded language's vdoc
+      const vdocHints = await withVirtualDocUri(vdoc, document.uri, "inlayHints", async (uri: Uri) => {
+        try {
+          return await commands.executeCommand<InlayHint[]>(
+            "vscode.executeInlayHintProvider",
+            uri,
+            vdocRange
+          );
+        } catch (error) {
+          const language = embeddedLanguage.ids.at(0) ?? "??";
+          console.warn(`[InlayHints] Error getting hints for language: ${language}. ${error}`);
+        }
+      })
+
+      if (!vdocHints) {
+        continue;
+      }
+
+      // Map results back to original doc range. Two places to update:
+      // - `InlayHint.position`
+      // - `InlayHint.textEdits.range`
+      vdocHints.forEach((hint) => {
+        // Unconditional `position` of where to show the inlay hint
+        hint.position = unadjustedPosition(vdoc.language, hint.position);
+
+        // Optional set of `textEdits` to "accept" the inlay hint
+        if (hint.textEdits) {
+          hint.textEdits.forEach((textEdit) => {
+            textEdit.range = unadjustedRange(vdoc.language, textEdit.range);
+          });
+        }
+      });
+
+      hints.push(...vdocHints);
+    }
+
+    return hints;
+  };
+}

apps/vscode/src/vdoc/vdoc.ts

@@ -114,6 +114,7 @@ export function virtualDocForCode(code: string[], language: EmbeddedLanguage) {
 export type VirtualDocAction =
   "completion" |
   "hover" |
+  "inlayHints" |
   "signature" |
   "definition" |
   "format" |
@@ -188,6 +189,39 @@ export function mainLanguage(
   }
 }
 
+/**
+ * Compute the unique set of `EmbeddedLanguages` found within this array of `Token[]`
+ */
+export function languages(
+  tokens: Token[]
+): EmbeddedLanguage[] {
+  const ids = new Set<string>();
+  const languages: EmbeddedLanguage[] = [];
+
+  tokens.filter(isExecutableLanguageBlock).forEach(token => {
+    const language = languageFromBlock(token);
+
+    if (!language) {
+      // This would be unexpected since we filtered to executable language blocks
+      return;
+    }
+
+    // `EmbeddedLanguage` doesn't have a unique identifier, it probably should.
+    // For now we join all `ids` and call that the identifier.
+    const id = language.ids.join("-");
+
+    if (ids.has(id)) {
+      // We've seen this `EmbeddedLanguage` already
+      return;
+    }
+
+    ids.add(id);
+    languages.push(language);
+  });
+
+  return languages;
+}
+
 export function languageFromBlock(token: Token) {
   const name = languageNameFromBlock(token);
   return embeddedLanguage(name);
@@ -211,6 +245,13 @@ export function unadjustedPosition(language: EmbeddedLanguage, pos: Position) {
   return new Position(pos.line - (language.inject?.length || 0), pos.character);
 }
 
+export function adjustedRange(language: EmbeddedLanguage, range: Range) {
+  return new Range(
+    adjustedPosition(language, range.start),
+    adjustedPosition(language, range.end)
+  );
+}
+
 export function unadjustedRange(language: EmbeddedLanguage, range: Range) {
   return new Range(
     unadjustedPosition(language, range.start),