vite-plugin-graphql-codegen-optimize

vitePluginGraphqlCodegenOptimize.ts

import {
  createSourceFile,
  forEachChild,
  isCallExpression,
  isIdentifier,
  isNoSubstitutionTemplateLiteral,
  isVariableStatement,
  visitNode,
  ScriptTarget,
  Node,
  CallExpression,
  SourceFile,
} from 'typescript'
import { parse as parseGql } from 'graphql/language/parser.js'
import { Plugin } from 'vite'

type TransformData = {
  operationNames: string[]
  replaces: { from: string; to: string }[]
}

const isFunctionCallExpression = (node: Node, functionName: string): node is CallExpression =>
  isCallExpression(node) && isIdentifier(node.expression) && node.expression.text === functionName

const inlineGraphqlCall = (node: CallExpression, sourceFile: SourceFile, transformData: TransformData) => {
  const [firstArg] = node.arguments
  if (!isNoSubstitutionTemplateLiteral(firstArg)) return

  const rawText = firstArg.rawText
  if (!rawText) return

  const ast = parseGql(rawText)
  const [firstDefinition] = ast.definitions
  if (firstDefinition.kind !== 'FragmentDefinition' && firstDefinition.kind !== 'OperationDefinition') return

  const name = firstDefinition.name?.value
  if (!name) return

  const isGraphqlCallResultAssignedToVariable = () => {
    let parentNode: undefined | Node = node.parent
    while (parentNode) {
      if (isVariableStatement(parentNode)) return true

      parentNode = parentNode.parent
    }

    return false
  }

  const graphqlCallExpression = node.getText(sourceFile)

  if (!isGraphqlCallResultAssignedToVariable()) {
    transformData.replaces.push({ from: graphqlCallExpression, to: '' })
    return
  }

  const operationName = firstDefinition.kind === 'OperationDefinition' ? `${name}Document` : `${name}FragmentDoc`
  transformData.replaces.push({ from: graphqlCallExpression, to: operationName })
  transformData.operationNames.push(operationName)
}

const inlineUseFragmentCall = (node: CallExpression, sourceFile: SourceFile, transformData: TransformData) => {
  const [, secondArg] = node.arguments
  if (!secondArg) return

  const useFragmentCallExpression = node.getText(sourceFile)
  const useFragmentSecondArg = secondArg.getText(sourceFile)
  transformData.replaces.push({ from: useFragmentCallExpression, to: useFragmentSecondArg })
}

export type PluginOptions = {
  graphqlFunctionName?: string
  useFragmentFunctionName?: string
  artifactDirectory: string
}

export default function (options: PluginOptions): Plugin {
  const graphqlFunctionName = options.graphqlFunctionName || 'graphql'
  const useFragmentFunctionName = options.useFragmentFunctionName || 'useFragment'

  return {
    name: 'vite-plugin-graphql-codegen-optimize',
    enforce: 'pre',
    transform: (code, id) => {
      const isTsFile = (id.endsWith('.ts') || id.endsWith('.tsx')) && !id.includes('node_modules')
      const hasGraphqlCallExpression = () => code.includes(`${graphqlFunctionName}(`)
      const hasUseFragmentCallExpression = () => code.includes(`${useFragmentFunctionName}(`)
      const isTransformTarget = isTsFile && (hasGraphqlCallExpression() || hasUseFragmentCallExpression())
      if (!isTransformTarget) return

      const transformData: TransformData = { operationNames: [], replaces: [] }

      const sourceFile = createSourceFile('', code, ScriptTarget.Latest, true)
      const visitor = (node: Node): undefined => {
        if (isFunctionCallExpression(node, graphqlFunctionName)) {
          inlineGraphqlCall(node, sourceFile, transformData)
          return
        }

        if (isFunctionCallExpression(node, useFragmentFunctionName)) {
          inlineUseFragmentCall(node, sourceFile, transformData)
          return
        }

        forEachChild(node, visitor)
      }
      visitNode(sourceFile, visitor)

      let transformedCode = code
      if (transformData.operationNames.length) {
        transformedCode =
          `import { ${transformData.operationNames.join(', ')} } from '${options.artifactDirectory}'\n` +
          transformedCode
      }
      transformData.replaces.forEach(({ from, to }) => {
        transformedCode = transformedCode.replace(from, to)
      })

      return transformedCode
    },
  }
}