feat: add pretty formatting for groq queries [WIP]

package.json

@@ -150,5 +150,9 @@
     "env:source": "export $(cat .envrc | xargs)",
     "vsce:publish": "sh publish.sh",
     "upgrade-interactive": "npx npm-check -u"
-  }
+  },
+  "files": [
+    "out/**/*",
+    "vendor/**/*"
+  ]
 }

src/extension.ts

@@ -7,6 +7,7 @@ import {Config} from './config/findConfig'
 import {GroqContentProvider} from './providers/content-provider'
 import {GROQCodeLensProvider} from './providers/groq-codelens-provider'
 import {executeGroq} from './query'
+import {formatGroq} from './format'
 
 export function activate(context: vscode.ExtensionContext) {
   // needed to load sanity.cli.ts
@@ -88,6 +89,21 @@ export function activate(context: vscode.ExtensionContext) {
   })
   context.subscriptions.push(disposable)
 
+  context.subscriptions.push(
+    vscode.commands.registerCommand('sanity.formatGroq', async (groqQuery: string, uri: string, range: vscode.Range) => {
+      const document = await vscode.workspace.openTextDocument(vscode.Uri.file(uri));
+      const query = groqQuery;
+      const formattedQuery = await formatGroq(query); // Await the async function
+
+      const edit = new vscode.WorkspaceEdit();
+      edit.replace(document.uri, range, formattedQuery);
+      await vscode.workspace.applyEdit(edit);
+
+      // Refresh the CodeLens
+      await vscode.commands.executeCommand('vscode.executeCodeLensProvider', document.uri);
+    })
+  );
+
   function readConfig() {
     const settings = vscode.workspace.getConfiguration('sanity')
     openJSONFile = settings.get('openJSONFile', false)

src/format.ts

@@ -0,0 +1,15 @@
+import { loadWasm } from "./wasm-loader"
+
+let isWasmLoaded = false;
+
+async function formatGroq(groq: string): Promise<string> {
+  if (!isWasmLoaded) {
+    await loadWasm();
+    isWasmLoaded = true;
+  }
+
+  const formattedGroq = (global as any).groqfmt(groq); // Assuming groqfmt is exposed globally by the WASM module
+  return formattedGroq.result;
+}
+
+export { formatGroq };

src/providers/groq-codelens-provider.ts

@@ -3,77 +3,93 @@ import {type CodeLensProvider, type TextDocument, type CancellationToken, CodeLe
 interface ExtractedQuery {
   content: string
   uri: string
-  position: Position
+  range: Range;
 }
 
 function extractAllTemplateLiterals(document: TextDocument): ExtractedQuery[] {
-  const documents: ExtractedQuery[] = []
-  const text = document.getText()
-  const regExpGQL = new RegExp('groq\\s*`([\\s\\S]+?)`', 'mg')
+  const documents: ExtractedQuery[] = [];
+  const text = document.getText();
+  const regExpGQL = new RegExp('groq\\s*`([\\s\\S]+?)`', 'mg');
 
-  let prevIndex = 0
-  let result
+  let result;
   while ((result = regExpGQL.exec(text)) !== null) {
-    const content = result[1]
-    const queryPosition = text.indexOf(content, prevIndex)
+    const content = result[1];
+    const startPosition = document.positionAt(result.index + result[0].indexOf(content));
+    const endPosition = document.positionAt(result.index + result[0].indexOf(content) + content.length);
+    const range = new Range(startPosition, endPosition);
     documents.push({
       content: content,
       uri: document.uri.path,
-      position: document.positionAt(queryPosition),
-    })
-    prevIndex = queryPosition + 1
+      range: range,
+    });
   }
-  return documents
+  return documents;
 }
 
 function extractAllDefineQuery(document: TextDocument): ExtractedQuery[] {
-  const documents: ExtractedQuery[] = []
-  const text = document.getText()
-  const pattern = '(\\s*defineQuery\\((["\'`])([\\s\\S]*?)\\2\\))'
+  const documents: ExtractedQuery[] = [];
+  const text = document.getText();
+  const pattern = '(\\s*defineQuery\\((["\'`])([\\s\\S]*?)\\2\\))';
   const regexp = new RegExp(pattern, 'g');
 
-  let prevIndex = 0
-  let result
+  let result;
   while ((result = regexp.exec(text)) !== null) {
-    const content = result[3]
-    const queryPosition = text.indexOf(result[1], prevIndex)
+    const content = result[3];
+    const startPosition = document.positionAt(result.index + result[0].indexOf(content));
+    const endPosition = document.positionAt(result.index + result[0].indexOf(content) + content.length);
+    const range = new Range(startPosition, endPosition);
     documents.push({
       content: content,
       uri: document.uri.path,
-      position: document.positionAt(queryPosition),
-    })
-    prevIndex = queryPosition + 1
+      range: range,
+    });
   }
-  return documents
+  return documents;
 }
 
 export class GROQCodeLensProvider implements CodeLensProvider {
   constructor() {}
 
   public provideCodeLenses(document: TextDocument, _token: CancellationToken): CodeLens[] {
+
     if (document.languageId === 'groq') {
       return [
         new CodeLens(new Range(new Position(0, 0), new Position(0, 0)), {
           title: 'Execute Query',
           command: 'sanity.executeGroq',
           arguments: [document.getText()],
         }),
-      ]
+        new CodeLens(new Range(new Position(0, 0), new Position(0, 0)), {
+          title: 'Format Query',
+          command: 'sanity.formatGroq',
+          arguments: [document.getText(), document.uri, new Position(0, 0)],
+        }),
+      ];
     }
 
     // find all lines where "groq" exists
-    const queries: ExtractedQuery[] = [...extractAllTemplateLiterals(document), ...extractAllDefineQuery(document)]
+    const queries: ExtractedQuery[] = [...extractAllTemplateLiterals(document), ...extractAllDefineQuery(document)];
 
     // add a button above each line that has groq
-    return queries.map((def) => {
-      return new CodeLens(
-        new Range(new Position(def.position.line, 0), new Position(def.position.line, 0)),
+    const lenses: CodeLens[] = [];
+    queries.forEach((def) => {
+      lenses.push(new CodeLens(
+        def.range,
         {
           title: 'Execute Query',
           command: 'sanity.executeGroq',
           arguments: [def.content],
         }
-      )
-    })
+      ));
+      lenses.push(new CodeLens(
+        def.range,
+        {
+          title: 'Format Query',
+          command: 'sanity.formatGroq',
+          arguments: [def.content, def.uri, def.range],
+        }
+      ));
+    });
+    return lenses;
   }
 }

src/wasm-loader.ts

@@ -0,0 +1,17 @@
+import * as path from 'path';
+import * as fs from 'fs';
+import * as util from 'util';
+
+import '@groqfmt/wasm/dist/wasm-exec';
+
+const go = new (global as any).Go();
+
+async function loadWasm() {
+  const wasmPath = path.resolve(__dirname, '../vendor/groqfmt.wasm');
+
+  const wasmData = await util.promisify(fs.readFile)(wasmPath);
+  const { instance } = await WebAssembly.instantiate(wasmData, go.importObject);
+  go.run(instance);
+}
+
+export { loadWasm };

tsconfig.json

@@ -17,5 +17,9 @@
     "esModuleInterop": true,
     "jsx": "react"
   },
-  "exclude": ["node_modules", ".vscode-test", "ts-graphql-plugin"]
+  "exclude": ["node_modules", ".vscode-test", "ts-graphql-plugin"],
+  "include": [
+    "src/**/*.ts",
+    "vendor/groqfmt.wasm"
+  ]
 }