TypeScript codemod to add history prop drilling where it was missing

morph.ts

/* eslint-disable no-unused-expressions */
/* eslint-disable @typescript-eslint/prefer-includes */

import {
    Project,
    Diagnostic,
    SyntaxKind,
    Node,
    StructureKind,
    QuoteKind,
    VariableDeclaration,
    Identifier,
    ts,
    BindingElement,
    ImportDeclaration,
    DiagnosticMessageChain,
} from 'ts-morph'
import { anyOf, isDefined } from '../../shared/src/util/types'
import * as prettier from 'prettier'

const project = new Project({
    tsConfigFilePath: 'tsconfig.json',
    manipulationSettings: {
        quoteKind: QuoteKind.Single,
        useTrailingCommas: true,
    },
})
// project.enableLogging(true)
project.addSourceFilesFromTsConfig('web/tsconfig.json')
project.addSourceFilesFromTsConfig('shared/tsconfig.json')
project.addSourceFilesAtPaths(['web/src/**/*.d.ts', 'shared/src/**/*.d.ts'])

console.log('Getting diagnostics')
const diagnostics = project
    .getPreEmitDiagnostics()
    .filter(d => !/(declaration|type definition) file/i.test(project.formatDiagnosticsWithColorAndContext([d])))

async function main(): Promise<void> {
    for (const diagnostic of diagnostics.filter(d =>
        /property 'history' is missing/i.test(project.formatDiagnosticsWithColorAndContext([d]))
    )) {
        try {
            const sourceFile = diagnostic.getSourceFile()
            if (!sourceFile) {
                continue
            }
            console.log(sourceFile.getFilePath())
            const dNode = sourceFile.getDescendantAtPos(diagnostic.getStart()!)!
            const jsxNode = dNode.getFirstAncestorOrThrow(anyOf(Node.isJsxSelfClosingElement, Node.isJsxOpeningElement))
            let initializer: string
            if (sourceFile?.getBaseName().endsWith('.test.tsx')) {
                // Test files
                // Add import { createMemoryHistory } from 'history' if not exists
                let importDecl = sourceFile.getImportDeclaration(decl => decl.getModuleSpecifierValue() === 'history')
                if (!importDecl) {
                    importDecl = sourceFile.addImportDeclaration({
                        namedImports: ['createMemoryHistory'],
                        moduleSpecifier: 'history',
                    })
                }
                const namedBindings = importDecl.getImportClauseOrThrow().getNamedBindingsOrThrow()
                if (
                    Node.isNamedImports(namedBindings) &&
                    !namedBindings
                        .getChildrenOfKind(SyntaxKind.ImportSpecifier)
                        .some(spec => spec.getText() === 'createMemoryHistory')
                ) {
                    importDecl.addNamedImport('createMemoryHistory')
                }
                const namespace = Node.isNamespaceImport(namedBindings) ? namedBindings.getName() : null
                initializer = namespace ? `{${namespace}.createMemoryHistory()}` : '{createMemoryHistory()}'
            } else {
                // Application code
                // Add import * as H from 'history' if not exists
                let importDecl = sourceFile.getImportDeclaration(decl => decl.getModuleSpecifierValue() === 'history')
                if (!importDecl) {
                    importDecl = sourceFile.addImportDeclaration({
                        namespaceImport: 'H',
                        moduleSpecifier: 'history',
                    })
                }
                const defaultImport = importDecl.getImportClauseOrThrow().getDefaultImport()
                if (defaultImport) {
                    // history should not be imported with a default import
                    importDecl = importDecl.replaceWithText("import * as H from 'history'")
                }
                const namedBindings = importDecl.getImportClauseOrThrow().getNamedBindingsOrThrow()
                if (
                    Node.isNamedImports(namedBindings) &&
                    !namedBindings
                        .getChildrenOfKind(SyntaxKind.ImportSpecifier)
                        .some(spec => spec.getText() === 'History')
                ) {
                    importDecl.addNamedImport('History')
                }
                const namespace = Node.isNamespaceImport(namedBindings) ? namedBindings.getName() : null
                const classDecl = jsxNode.getFirstAncestor(anyOf(Node.isClassDeclaration, Node.isClassExpression))
                if (classDecl) {
                    // React class component
                    initializer = '{this.props.history}'
                    // Add history prop to Props interface
                    const [propsTypeArg] = classDecl.getExtendsOrThrow().getTypeArguments()
                    const propsSymbol = Node.isTypeReferenceNode(propsTypeArg)
                        ? propsTypeArg.getFirstChildByKindOrThrow(SyntaxKind.Identifier).getSymbolOrThrow()
                        : propsTypeArg.getSymbolOrThrow()
                    if (!propsSymbol.getMember('history')) {
                        const propsDecl = propsSymbol.getDeclarations()[0]
                        if (Node.isInterfaceDeclaration(propsDecl) || Node.isTypeLiteralNode(propsDecl)) {
                            propsDecl.addProperty({
                                name: 'history',
                                type: namespace ? `${namespace}.History` : 'History',
                            })
                        } else {
                            throw new Error('Props type is neither interface nor type literal')
                        }
                    }
                } else {
                    // Function component
                    const functionDecl = jsxNode.getFirstAncestorOrThrow((node: Node): node is VariableDeclaration => {
                        const isVarDecl = Node.isVariableDeclaration(node)
                        try {
                            const isFuncComp = node.getType().getSymbol()?.getName() === 'FunctionComponent'
                            return isVarDecl && isFuncComp
                        } catch {
                            return false
                        }
                    })
                    const propsSymbol = functionDecl.getType().getTypeArguments()[0].getSymbolOrThrow()
                    if (!propsSymbol.getMember('history')) {
                        const propsDecl = propsSymbol.getDeclarations()[0]
                        if (Node.isInterfaceDeclaration(propsDecl) || Node.isTypeLiteralNode(propsDecl)) {
                            propsDecl.addProperty({
                                name: 'history',
                                type: namespace ? `${namespace}.History` : 'History',
                            })
                        } else {
                            throw new Error('Props type is neither interface nor type literal')
                        }
                    }
                    const paramNode = functionDecl.getFirstDescendantByKindOrThrow(SyntaxKind.Parameter)
                    const paramName = paramNode.getNameNode()
                    if (Node.isObjectBindingPattern(paramName)) {
                        const restSpread = paramName.getFirstChild(
                            (node): node is BindingElement => Node.isBindingElement(node) && !!node.getDotDotDotToken()
                        )
                        if (restSpread) {
                            // Reference rest spread
                            initializer = `{${restSpread.getName()}.history}`
                        } else {
                            if (!paramName.getElements().some(element => element.getName() === 'history')) {
                                // Add to destructured props
                                paramName.transform(traversal => {
                                    if (ts.isObjectBindingPattern(traversal.currentNode)) {
                                        return ts.updateObjectBindingPattern(traversal.currentNode, [
                                            ...traversal.currentNode.elements,
                                            ts.createBindingElement(undefined, undefined, 'history'),
                                        ])
                                    }
                                    return traversal.currentNode
                                })
                            }
                            initializer = '{history}'
                        }
                    } else {
                        initializer = '{' + (paramName as Identifier).getText() + '.history}'
                    }
                }
            }

            jsxNode.addAttribute({
                kind: StructureKind.JsxAttribute,
                name: 'history',
                initializer,
            })
            sourceFile.replaceWithText(
                prettier.format(sourceFile.getFullText(), {
                    ...(await prettier.resolveConfig(sourceFile.getFilePath()))!,
                    filepath: sourceFile.getFilePath(),
                })
            )
            await sourceFile.save()
        } catch (err) {
            console.error(err)
        }
    }
}
// eslint-disable-next-line @typescript-eslint/no-floating-promises
main()