From c518591250575a8a9aa5e045c4cb589ad796e5aa Mon Sep 17 00:00:00 2001
From: Owen Kieffer-Jones <owen.kieffer-jones@outletcity.com>
Date: Tue, 5 Nov 2024 09:59:12 +0100
Subject: [PATCH] feat(codegen): fix literal function arguments and add default
 values

---
 .../__tests__/findQueriesInSource.test.ts     | 15 +++++++++
 .../src/typescript/expressionResolvers.ts     | 32 +++++++++++++++----
 2 files changed, 41 insertions(+), 6 deletions(-)

diff --git a/packages/@sanity/codegen/src/typescript/__tests__/findQueriesInSource.test.ts b/packages/@sanity/codegen/src/typescript/__tests__/findQueriesInSource.test.ts
index 82a75657a8b..0b712a5c0a9 100644
--- a/packages/@sanity/codegen/src/typescript/__tests__/findQueriesInSource.test.ts
+++ b/packages/@sanity/codegen/src/typescript/__tests__/findQueriesInSource.test.ts
@@ -63,6 +63,21 @@ describe('findQueries with the groq template', () => {
       expect(queryResult?.result).toEqual('*[_type == "author"]')
     })
 
+    test('with function with literal parameters', () => {
+      const source = `
+      import { groq } from "groq";
+      const getType = (type, x = '2') => () =>  \`\${type}\${x}\`;
+      const query = groq\`*[_type == "\${getType("author")()}"]\`
+      const res = sanity.fetch(query);
+    `
+
+      const queries = findQueriesInSource(source, 'test.ts')
+
+      const queryResult = queries[0]
+
+      expect(queryResult?.result).toEqual('*[_type == "author2"]')
+    })
+
     test('with block comment', () => {
       const source = `
         import { groq } from "groq";
diff --git a/packages/@sanity/codegen/src/typescript/expressionResolvers.ts b/packages/@sanity/codegen/src/typescript/expressionResolvers.ts
index 468f0d14d15..8e30fa66b78 100644
--- a/packages/@sanity/codegen/src/typescript/expressionResolvers.ts
+++ b/packages/@sanity/codegen/src/typescript/expressionResolvers.ts
@@ -2,7 +2,7 @@ import fs from 'node:fs'
 import path from 'node:path'
 
 import {type TransformOptions} from '@babel/core'
-import traverse, {type Scope} from '@babel/traverse'
+import traverse, {Scope} from '@babel/traverse'
 import * as babelTypes from '@babel/types'
 import createDebug from 'debug'
 
@@ -125,12 +125,13 @@ export function resolveExpression({
   }
 
   if (babelTypes.isVariableDeclarator(node)) {
-    if (!node.init) {
+    const init = node.init ?? (babelTypes.isAssignmentPattern(node.id) && node.id.right)
+    if (!init) {
       throw new Error(`Unsupported variable declarator`)
     }
 
     return resolveExpression({
-      node: node.init,
+      node: init,
       fnArguments,
       scope,
       filename,
@@ -174,11 +175,20 @@ export function resolveExpression({
     babelTypes.isFunctionDeclaration(node) ||
     babelTypes.isFunctionExpression(node)
   ) {
+    const newScope = new Scope(scope.path, scope)
+
+    params.forEach((param, i) => {
+      newScope.push({
+        id: param as babelTypes.LVal,
+        init: fnArguments[i] as babelTypes.Expression | undefined,
+      })
+    })
+
     return resolveExpression({
       node: node.body,
       params: node.params,
       fnArguments,
-      scope,
+      scope: newScope,
       filename,
       file,
       babelConfig,
@@ -198,7 +208,7 @@ export function resolveExpression({
   }
 
   if (babelTypes.isImportDefaultSpecifier(node) || babelTypes.isImportSpecifier(node)) {
-    return resolveImportSpecifier({node, file, scope, filename, resolver, babelConfig})
+    return resolveImportSpecifier({node, file, scope, filename, fnArguments, resolver, babelConfig})
   }
 
   if (babelTypes.isAssignmentPattern(node)) {
@@ -245,7 +255,10 @@ function resolveIdentifier({
         babelTypes.isIdentifier(param.left) &&
         node.name === param.left.name),
   )
-  const argument = fnArguments[paramIndex]
+  let argument = fnArguments[paramIndex]
+  if (!argument && paramIndex >= 0 && babelTypes.isAssignmentPattern(params[paramIndex])) {
+    argument = params[paramIndex].right
+  }
   if (argument && babelTypes.isLiteral(argument)) {
     return resolveExpression({
       node: argument,
@@ -320,6 +333,7 @@ function resolveImportSpecifier({
   node,
   file,
   filename,
+  fnArguments,
   resolver,
   babelConfig,
 }: {
@@ -327,6 +341,7 @@ function resolveImportSpecifier({
   file: babelTypes.File
   scope: Scope
   filename: string
+  fnArguments: babelTypes.Node[]
   resolver: NodeJS.RequireResolve
   babelConfig: TransformOptions
 }): resolveExpressionReturnType {
@@ -382,6 +397,7 @@ function resolveImportSpecifier({
       node: binding.path.node,
       file: tree,
       scope: newScope,
+      fnArguments,
       babelConfig,
       filename: resolvedFile,
       resolver,
@@ -413,6 +429,7 @@ function resolveImportSpecifier({
       node: namedExport,
       importName: newImportName,
       filename: resolvedFile,
+      fnArguments,
       resolver,
       babelConfig,
     })
@@ -446,12 +463,14 @@ function resolveExportSpecifier({
   node,
   importName,
   filename,
+  fnArguments,
   babelConfig,
   resolver,
 }: {
   node: babelTypes.ExportNamedDeclaration | babelTypes.ExportAllDeclaration
   importName: string
   filename: string
+  fnArguments: babelTypes.Node[]
   babelConfig: TransformOptions
   resolver: NodeJS.RequireResolve
 }): resolveExpressionReturnType {
@@ -484,6 +503,7 @@ function resolveExportSpecifier({
       filename: importFileName,
       babelConfig,
       resolver,
+      fnArguments,
     })
   }