kea-typegen 3.6.6 → 3.7.0

package/src/types.ts

@@ -82,6 +82,8 @@ export interface AppOptions {
     showTsErrors?: boolean
     /** Cache generated logic files into .typegen, use them if generating a logic type for the first time */
     useCache?: boolean
+    /** Skip Prettier formatting while generating logic type files */
+    prettier?: boolean
 
     log: (message: string) => void
 }

package/src/visit/visit.ts

@@ -25,7 +25,7 @@ import { visitEvents } from './visitEvents'
 import { visitDefaults } from './visitDefaults'
 import { visitSharedListeners } from './visitSharedListeners'
 import { cloneNode } from 'ts-clone-node'
-import {NodeBuilderFlags} from "typescript";
+import { NodeBuilderFlags } from 'typescript'
 
 const visitFunctions = {
     actions: visitActions,
@@ -44,20 +44,52 @@ const visitFunctions = {
     windowValues: visitWindowValues,
 }
 
+// Cheap text prefilters so we only do expensive AST + type-checker work on files
+// that can actually contain the thing we're looking for. On large projects most
+// source files contain no Kea logic at all, and walking every node (and resolving
+// a symbol for every identifier via the checker) for them dominates runtime.
+//
+// `kea` is never aliased on import in practice, so a `kea(` / `kea<` text match is a
+// safe-enough gate for the Kea-call pass. If a project does alias `kea`, set
+// KEA_TYPEGEN_NO_PREFILTER=1 to fall back to scanning every file.
+const KEA_CALL_REGEX = /\bkea\s*[<(]/
+const prefilterDisabled = process.env.KEA_TYPEGEN_NO_PREFILTER === '1'
+
+function fileMightContainKeaCall(sourceFile: ts.SourceFile): boolean {
+    return prefilterDisabled || KEA_CALL_REGEX.test(sourceFile.text)
+}
+
+function fileMightContainResetContext(sourceFile: ts.SourceFile): boolean {
+    return prefilterDisabled || sourceFile.text.includes('resetContext')
+}
+
 export function visitProgram(program: ts.Program, appOptions?: AppOptions): ParsedLogic[] {
     const checker = program.getTypeChecker()
     const parsedLogics: ParsedLogic[] = []
     const pluginModules: PluginModule[] = []
     const typeBuilderModules: TypeBuilderModule[] = []
+    const sourceFilePath = appOptions?.sourceFilePath ? path.resolve(appOptions.sourceFilePath) : undefined
 
     for (const sourceFile of program.getSourceFiles()) {
-        if (!sourceFile.isDeclarationFile && !sourceFile.fileName.endsWith('Type.ts')) {
+        if (
+            !sourceFile.isDeclarationFile &&
+            !sourceFile.fileName.endsWith('Type.ts') &&
+            fileMightContainResetContext(sourceFile)
+        ) {
             ts.forEachChild(sourceFile, visitResetContext(checker, pluginModules))
         }
     }
 
     for (const sourceFile of program.getSourceFiles()) {
-        if (!sourceFile.isDeclarationFile && !sourceFile.fileName.endsWith('Type.ts')) {
+        if (sourceFilePath && path.resolve(sourceFile.fileName) !== sourceFilePath) {
+            continue
+        }
+
+        if (
+            !sourceFile.isDeclarationFile &&
+            !sourceFile.fileName.endsWith('Type.ts') &&
+            fileMightContainKeaCall(sourceFile)
+        ) {
             if (appOptions?.verbose) {
                 appOptions.log(`👀 Visiting: ${path.relative(process.cwd(), sourceFile.fileName)}`)
             }
@@ -245,7 +277,6 @@ export function visitKeaCalls(
         const calls: {
             name: string
             type: ts.Type
-            typeNode: ts.TypeNode | null
             expression: ts.Expression
             typeBuilders: PluginModule[]
         }[] = []
@@ -261,8 +292,7 @@ export function visitKeaCalls(
                 }
                 const name = symbol.getName()
                 const type = checker.getTypeOfSymbolAtLocation(symbol, symbol.valueDeclaration)
-                const typeNode = type ? checker.typeToTypeNode(type, undefined, NodeBuilderFlags.NoTruncation) : null
-                calls.push({ name, type, typeNode, expression: inputProperty.initializer, typeBuilders: [] })
+                calls.push({ name, type, expression: inputProperty.initializer, typeBuilders: [] })
             }
         } else if (ts.isArrayLiteralExpression(input)) {
             for (const callExpression of input.elements) {
@@ -273,7 +303,6 @@ export function visitKeaCalls(
                 const builderName = callExpression.expression.getText()
                 const argument = callExpression.arguments[0]
                 const type = checker.getTypeAtLocation(argument)
-                const typeNode = type ? checker.typeToTypeNode(type, undefined, NodeBuilderFlags.NoTruncation) : null
 
                 const identifier = callExpression.expression
                 const symbol = checker.getSymbolAtLocation(identifier)
@@ -348,14 +377,13 @@ export function visitKeaCalls(
                 calls.push({
                     name: builderName,
                     type,
-                    typeNode,
                     expression: callExpression.arguments[0],
                     typeBuilders: typeBuilders,
                 })
             }
         }
 
-        for (let { name, type, typeNode, expression, typeBuilders } of calls) {
+        for (let { name, type, expression, typeBuilders } of calls) {
             if (name === 'path' || name === 'logicPath') {
                 parsedLogic.hasPathInLogic = true
             }
@@ -363,18 +391,31 @@ export function visitKeaCalls(
                 parsedLogic.hasKeyInLogic = true
             }
 
-            if (typeNode && ts.isFunctionTypeNode(typeNode)) {
-                type = type.getCallSignatures()[0].getReturnType()
-                typeNode = type ? checker.typeToTypeNode(type, undefined, NodeBuilderFlags.NoTruncation) : null
+            // Equivalent to the previous `ts.isFunctionTypeNode(typeNode)` check, but without
+            // eagerly building a typeNode: a pure function type has call signatures and no
+            // members, so we unwrap it to its return type before visiting/printing.
+            const callSignatures = type.getCallSignatures()
+            if (callSignatures.length > 0 && type.getProperties().length === 0) {
+                type = callSignatures[0].getReturnType()
             }
 
             visitFunctions[name]?.(parsedLogic, type, expression)
 
+            let typeNode: ts.TypeNode | null | undefined = undefined
+            const getTypeNode = (): ts.TypeNode | null => {
+                if (typeNode === undefined) {
+                    typeNode = type ? checker.typeToTypeNode(type, undefined, NodeBuilderFlags.NoTruncation) : null
+                }
+                return typeNode
+            }
+
             const visitKeaPropertyArguments: VisitKeaPropertyArguments = {
                 name,
                 appOptions,
                 type,
-                typeNode,
+                get typeNode() {
+                    return getTypeNode()
+                },
                 parsedLogic,
                 node: expression,
                 checker,

package/src/write/writeTypeImports.ts

@@ -1,10 +1,41 @@
 import { AppOptions, ParsedLogic } from '../types'
 import * as ts from 'typescript'
-import { print, visit } from 'recast'
 import * as osPath from 'path'
 import { runThroughPrettier } from '../print/print'
 import * as fs from 'fs'
-import { t, b, visitAllKeaCalls, getAst } from './utils'
+interface TextEdit {
+    start: number
+    end: number
+    text: string
+}
+
+function applyTextEdits(source: string, edits: TextEdit[]): string {
+    return edits
+        .sort((a, b) => b.start - a.start || b.end - a.end)
+        .reduce((output, { start, end, text }) => output.slice(0, start) + text + output.slice(end), source)
+}
+
+function getImportPath(importDeclaration: ts.ImportDeclaration): string | null {
+    return ts.isStringLiteralLike(importDeclaration.moduleSpecifier) ? importDeclaration.moduleSpecifier.text : null
+}
+
+function getImportInsertPosition(sourceFile: ts.SourceFile, rawCode: string): number {
+    const importDeclarations = sourceFile.statements.filter(ts.isImportDeclaration)
+    if (importDeclarations.length > 0) {
+        return importDeclarations[importDeclarations.length - 1].getEnd()
+    }
+
+    const shebangMatch = rawCode.match(/^#!.*(?:\r?\n|$)/)
+    return shebangMatch ? shebangMatch[0].length : 0
+}
+
+function getTypeArgumentInsertEnd(callExpression: ts.CallExpression, sourceFile: ts.SourceFile): number {
+    const openParenToken = callExpression
+        .getChildren(sourceFile)
+        .find((child) => child.kind === ts.SyntaxKind.OpenParenToken)
+
+    return openParenToken ? openParenToken.getStart(sourceFile) : callExpression.expression.getEnd()
+}
 
 export async function writeTypeImports(
     appOptions: AppOptions,
@@ -15,9 +46,10 @@ export async function writeTypeImports(
 ) {
     const { log } = appOptions
     const sourceFile = program.getSourceFile(filename)
-    const rawCode = sourceFile.getText()
-
-    const ast = getAst(filename, rawCode)
+    if (!sourceFile) {
+        throw new Error(`Could not find source file: ${filename}`)
+    }
+    const rawCode = fs.readFileSync(filename, 'utf8')
 
     let importLocation = osPath
         .relative(osPath.dirname(filename), logicsNeedingImports[0].typeFileName)
@@ -26,56 +58,57 @@ export async function writeTypeImports(
         importLocation = `./${importLocation}`
     }
 
-    // add import if missing
-    let foundImport = false
-    visit(ast, {
-        visitImportDeclaration(path) {
-            const importPath =
-                path.value.source && t.StringLiteral.check(path.value.source) ? path.value.source.value : null
-
-            if (
-                t.ImportDeclaration.check(path.value) &&
-                importPath &&
-                osPath.resolve(osPath.dirname(filename), importPath) ===
-                    osPath.resolve(osPath.dirname(filename), importLocation)
-            ) {
-                foundImport = true
-                path.value.importKind = 'type'
-                path.value.specifiers = parsedLogics.map((l) =>
-                    b.importSpecifier(b.identifier(l.logicTypeName), b.identifier(l.logicTypeName)),
-                )
-            }
-            return false
-        },
+    const desiredImport = `import type { ${parsedLogics.map((l) => l.logicTypeName).join(', ')} } from '${importLocation}'`
+    const importDeclarations = sourceFile.statements.filter(ts.isImportDeclaration)
+    const matchingImport = importDeclarations.find((importDeclaration) => {
+        const importPath = getImportPath(importDeclaration)
+        return (
+            importPath !== null &&
+            osPath.resolve(osPath.dirname(filename), importPath) ===
+                osPath.resolve(osPath.dirname(filename), importLocation)
+        )
     })
 
-    if (!foundImport) {
-        visit(ast, {
-            visitProgram(path) {
-                path.value.body = [
-                    ...path.value.body.filter((n) => t.ImportDeclaration.check(n)),
-                    b.importDeclaration(
-                        parsedLogics.map((l) =>
-                            b.importSpecifier(b.identifier(l.logicTypeName), b.identifier(l.logicTypeName)),
-                        ),
-                        b.stringLiteral(importLocation),
-                        'type',
-                    ),
-                    ...path.value.body.filter((n) => !t.ImportDeclaration.check(n)),
-                ]
-                return false
-            },
+    const edits: TextEdit[] = []
+    if (matchingImport) {
+        edits.push({
+            start: matchingImport.getStart(sourceFile),
+            end: matchingImport.getEnd(),
+            text: desiredImport,
+        })
+    } else {
+        const insertPos = getImportInsertPosition(sourceFile, rawCode)
+        const hasExistingImports = importDeclarations.length > 0
+        const importText = hasExistingImports
+            ? `\n${desiredImport}${rawCode.slice(insertPos, insertPos + 1) === '\n' ? '' : '\n'}`
+            : `${desiredImport}\n`
+
+        edits.push({
+            start: insertPos,
+            end: insertPos,
+            text: importText,
         })
     }
 
-    // find all kea calls, add `<logicType>` type parameters if needed
-    visitAllKeaCalls(ast, logicsNeedingImports, filename, ({ path, parsedLogic }) => {
-        path.node.typeParameters = b.tsTypeParameterInstantiation([
-            b.tsTypeReference(b.identifier(parsedLogic.logicTypeName)),
-        ])
-    })
+    for (const parsedLogic of logicsNeedingImports) {
+        const callExpression = parsedLogic.node.parent
+        if (!ts.isCallExpression(callExpression)) {
+            continue
+        }
+
+        const typeArgumentStart = callExpression.expression.getEnd()
+        const typeArgumentEnd = callExpression.typeArguments
+            ? getTypeArgumentInsertEnd(callExpression, sourceFile)
+            : typeArgumentStart
+
+        edits.push({
+            start: typeArgumentStart,
+            end: typeArgumentEnd,
+            text: `<${parsedLogic.logicTypeName}>`,
+        })
+    }
 
-    const newText = await runThroughPrettier(print(ast).code, filename)
+    const newText = await runThroughPrettier(applyTextEdits(rawCode, edits), filename)
     fs.writeFileSync(filename, newText)
 
     log(`🔥 Import added: ${osPath.relative(process.cwd(), filename)}`)