convert stack utility functions to components using TS Morph

codemod.ts

#!/usr/bin/env node
import {
  ReferencedSymbol,
  Project,
  ts,
  Node,
  SourceFile,
  ImportDeclaration,
  ImportClause,
  JsxElement,
  JsxSelfClosingElement,
  JsxAttributeStructure,
  OptionalKind,
} from 'ts-morph'
import { resolve } from 'node:path'

// Utilities
function getImportDeclarationByModuleSpecifier(
  sourceFile: SourceFile,
  moduleSpecifier: string
): ImportDeclaration | undefined {
  return sourceFile.getImportDeclarations().find((importDeclaration) => {
    return importDeclaration.getModuleSpecifierValue() === moduleSpecifier
  })
}

function getImportClause(
  sourceFile: SourceFile,
  moduleSpecifier: string,
  importClause: string
): ImportClause | undefined {
  const importDeclaration = getImportDeclarationByModuleSpecifier(
    sourceFile,
    moduleSpecifier
  )

  if (!importDeclaration) {
    return
  }

  const importDeclarationClause = importDeclaration.getImportClause()

  if (importDeclarationClause.getDefaultImport().getText() === importClause) {
    return importDeclarationClause
  }
}

// Rename JSX Element Identifier accounting for opening, closing, and self closing elements
function renameJSXIdentifier(
  jsxElement: JsxElement | JsxSelfClosingElement,
  identifier: string
) {
  if (Node.isJsxSelfClosingElement(jsxElement)) {
    jsxElement
      .getFirstDescendantByKind(ts.SyntaxKind.Identifier)
      .rename(identifier)

    return
  }

  jsxElement
    .getOpeningElement()
    .getFirstDescendantByKind(ts.SyntaxKind.Identifier)
    .rename(identifier)

  jsxElement
    .getClosingElement()
    .getFirstDescendantByKind(ts.SyntaxKind.Identifier)
    .rename(identifier)
}

const sourceFilePath = process.argv[2] || '.'

const project = new Project({
  tsConfigFilePath: resolve(process.cwd(), sourceFilePath, 'tsconfig.json'),
})

let references: ReferencedSymbol[] = []

for (const sourceFile of project.getSourceFiles()) {
  for (const importDeclaration of sourceFile.getImportDeclarations()) {
    if (importDeclaration.getModuleSpecifierValue() === '@vercel/ui') {
      for (const namedImport of importDeclaration.getNamedImports()) {
        if (namedImport.getName() === 'stack') {
          const identifier = namedImport.getFirstDescendantByKindOrThrow(
            ts.SyntaxKind.Identifier
          )

          references = identifier.findReferences()

          break
        }
      }
    }
  }
}

const flattenedReferences = references
  .flatMap((reference) => reference.getReferences())
  .map((reference) => reference.getNode())

flattenedReferences.forEach((node) => {
  if (node.getParent().getKind() === ts.SyntaxKind.ImportSpecifier) {
    const importDeclaration = node.getFirstAncestorByKindOrThrow(
      ts.SyntaxKind.ImportDeclaration
    )
    const importSpecifierNode = node.getParent()
    const hasStackImport = getImportClause(
      node.getSourceFile(),
      '#/geist/stack',
      'Stack'
    )

    if (!hasStackImport) {
      node
        .getSourceFile()
        .insertImportDeclaration(importDeclaration.getChildIndex() + 1, {
          defaultImport: 'Stack',
          moduleSpecifier: '#/geist/stack',
        })
    }

    if (importDeclaration.getNamedImports().length === 1) {
      importDeclaration.remove()
    } else if (Node.isImportSpecifier(importSpecifierNode)) {
      importSpecifierNode.remove()
    }

    return
  }

  // Skip node module references to original package
  if (node.getSourceFile().isInNodeModules()) {
    return
  }

  const jsxElement = node.getFirstAncestorByKindOrThrow(
    ts.SyntaxKind.JsxElement
  )
  const jsxElementName = jsxElement
    .getOpeningElement()
    .getTagNameNode()
    .getText()

  // Replace opening/closing element identifiers with Stack
  renameJSXIdentifier(jsxElement, 'Stack')

  // Move object literal expression to attributes
  const objectLiteralExpression = node
    .getFirstAncestorByKindOrThrow(ts.SyntaxKind.CallExpression)
    .getFirstDescendantByKindOrThrow(ts.SyntaxKind.ObjectLiteralExpression)
  const attributesToAdd: OptionalKind<JsxAttributeStructure>[] = []
  const validAttributes = ['gap', 'direction', 'align', 'justify']
  const styleProperties = []

  if (jsxElementName !== 'div') {
    attributesToAdd.push({
      name: 'as',
      initializer: `"${jsxElementName}"`,
    })
  }

  objectLiteralExpression.getProperties().forEach((property) => {
    if (Node.isPropertyAssignment(property)) {
      const initializer = property.getInitializer()
      const initializerText = initializer.getText()
      const maybeNumber = parseInt(initializerText.slice(1, -1))
      let jsxInitializer = isNaN(maybeNumber)
        ? initializerText
        : `{${maybeNumber}}`
      const name = property
        .getFirstDescendantByKindOrThrow(ts.SyntaxKind.Identifier)
        .getText()

      if (name === 'gap') {
        jsxInitializer = `{${maybeNumber / 4}}`
      }

      if (initializer.getKind() === ts.SyntaxKind.ObjectLiteralExpression) {
        jsxInitializer = `{${initializerText}}`
      }

      if (validAttributes.includes(name)) {
        attributesToAdd.push({ name, initializer: jsxInitializer })
      } else {
        // Set margin and padding to number since it will be converted to a pixel value in React
        if (name.includes('margin') || name.includes('padding')) {
          property.setInitializer(String(maybeNumber))
        }

        styleProperties.push(property)
      }
    }
  })

  // Sort attributes by name
  attributesToAdd.sort((a, b) => {
    if (a.name < b.name) {
      return -1
    }

    if (a.name > b.name) {
      return 1
    }

    return 0
  })

  // Add attributes to opening element
  for (const attribute of attributesToAdd) {
    jsxElement.getOpeningElement().addAttribute(attribute)
  }

  if (styleProperties.length > 0) {
    jsxElement.getOpeningElement().addAttribute({
      name: 'style',
      initializer: `{{${styleProperties.map((p) => p.getText()).join(', ')}}}`,
    })
  }

  // Remove original attribute
  node.getFirstAncestorByKindOrThrow(ts.SyntaxKind.JsxAttribute).remove()
})

project.save()