fix: transformation of index accessed functions wraps code correctly

Author: skoshx

src/index.ts

@@ -20,24 +20,70 @@ import { getLoadedConfig } from './core/config'
 import { CaptureInvocation } from './exports'
 import { createHumanLog } from './core/errors'
 
-export function uff<T extends AnyFunction>(
+/* export function ufc<T extends AnyFunction>(
 	func: T,
-	// context: any = null,
-	id: string | null = null
-): T {
+	context: any = null
+) {
 	return function (this: any) {
+		// @ts-ignore
+		return func.apply(context ?? this, arguments)
+	} as T;
+} */
+export function ufc<T extends AnyFunction>(
+	functionOrNamespace: T,
+	opts: FlytrapCallOptions
+): ReturnType<T> {
+	try {
+		saveFunctionCall(opts)
+
+		const execFunctionCall = <T>(fn: T, name: string, args: any[]) => {
+			log.info('call-execution', `Executing function ${name}()`, { args })
+			// @ts-ignore
+			if (fn?.[name]) {
+				if (typeof fn === 'object' && typeof (fn as any)?.[name] === 'function') {
+					// @ts-ignore
+					return fn[name].call(fn, ...args)
+				}
+
+				// @ts-ignore
+				return fn[name](...args)
+			} else {
+				// @ts-ignore
+				return fn(...args)
+			}
+		}
+
+		const output = execFunctionCall(functionOrNamespace, opts.name, opts.args)
+		saveOutputForFunctionCall(opts.id, output)
+		return output
+	} catch (error) {
+		/**
+		 * Oops! We found a bug, let's send the current
+		 * executing function along with its data to the
+		 * Flytrap API.
+		 */
+		saveErrorForFunctionCall(opts.id, error as Error)
+		log.info('capture', `Captured error in function call with ID "${opts.id}".`, { error })
+		throw error
+	}
+}
+
+export function uff<T extends AnyFunction>(func: T, id: string | null = null): T {
+	return function (this: any, ...args: any[]) {
 		if (id) {
 			addFunctionInvocation(id, {
 				// eslint-disable-next-line
-				args: Array.from(arguments),
+				// args: Array.from(arguments),
+				args,
 				timestamp: Date.now()
 			})
 		}
 		try {
 			const context = null
 			// @ts-ignore
 			// eslint-disable-next-line
-			const functionOutput = func.apply(context ?? this, arguments)
+			// const functionOutput = func.apply(context ?? this, arguments)
+			const functionOutput = func.apply(context ?? this, args)
 			if (id) {
 				saveOutputForFunction(id, functionOutput)
 			}

src/transform/imports.ts

@@ -7,8 +7,7 @@ import { parseCode } from './parser'
 import { extname } from '../core/util'
 
 export function getRequiredExportsForCapture(): string[] {
-	// return ['useFlytrapCall', 'useFlytrapCallAsync', 'useFlytrapFunction', 'setFlytrapConfig', 'uff']
-	return ['uff']
+	return ['uff', 'ufc', 'setFlytrapConfig']
 }
 
 export function getCoreExports(): string[] {

src/transform/index.ts

@@ -17,9 +17,13 @@ import {
 	objectExpression,
 	V8IntrinsicIdentifier,
 	ObjectProperty,
-	VariableDeclarator
+	VariableDeclarator,
+	MemberExpression,
+	isMemberExpression,
+	isNumericLiteral,
+	isStringLiteral,
+	Identifier
 } from '@babel/types'
-import { parse } from '@babel/parser'
 import generate from '@babel/generator'
 
 import { getCoreExports } from './imports'
@@ -37,11 +41,36 @@ import { _babelInterop } from './util'
 import { parseCode } from './parser'
 import { createHumanLog } from '../core/errors'
 
+export function getCalleeAndAccessorKey(node: MemberExpression | Identifier) {
+	if (!isMemberExpression(node)) {
+		return {
+			callee: node,
+			accessorKey: undefined
+		}
+	}
+
+	// @ts-expect-error
+	const callee: MemberExpression | Identifier = node.object
+
+	let accessorKey = node.property
+
+	if (!node.computed) {
+		if (isIdentifier(node.property)) {
+			accessorKey = stringLiteral(node.property.name)
+		}
+		if (isNumericLiteral(node.property) || isStringLiteral(node.property)) {
+			accessorKey = stringLiteral(String(node.property.value))
+		}
+	}
+
+	return { callee, accessorKey }
+}
+
 export function shouldBeWrappedUff(path: NodePath) {
 	if (
 		isCallExpression(path.parent) &&
 		isIdentifier(path.parent.callee) &&
-		['uff'].includes(path.parent.callee.name)
+		['uff', 'ufc'].includes(path.parent.callee.name)
 	) {
 		return false
 	}
@@ -98,10 +127,6 @@ export function flytrapTransformUff(
 		? findIgnoredImports(code, config.packageIgnores)
 		: undefined
 
-	// What gets transformed and captured is defined in the config.
-	// const shouldTransformFunctions = config?.captureDataFrom === "calls" ? false : true;
-	// const shouldTransformCalls = config?.captureDataFrom === 'functions' ? false : true;
-
 	try {
 		_babelInterop(babelTraverse)(ast, {
 			...(!config?.transformOptions?.disableTransformation?.includes('arrow-function') && {
@@ -188,30 +213,40 @@ export function flytrapTransformUff(
 						return
 					}
 
-					const functionCallName = extractFunctionCallName(path.node)
+					const fullFunctionCallName = _babelInterop(generate)({
+						...path.node,
+						arguments: []
+					}).code.replaceAll('()', '')
+
+					if (fullFunctionCallName === 'this' || fullFunctionCallName.split('.').at(0) === 'this') {
+						return
+					}
+					const functionCallName = fullFunctionCallName.split('.').at(-1)!
 					const scopes = extractCurrentScope(path)
 					const functionCallId = extractFunctionCallId(path, filePath, functionCallName, scopes)
-					const useFunctionName = isAwaitExpression(path.parent)
-						? 'useFlytrapCallAsync'
-						: 'useFlytrapCall'
-
-					function transformCallee(callee: V8IntrinsicIdentifier | Expression) {
-						if (callee.type === 'MemberExpression') {
-							return callee.object
-						}
-						return callee
+					const useFunctionName = isAwaitExpression(path.parent) ? 'ufc' : 'ufc'
+
+					// @ts-ignore
+					const { callee, accessorKey } = getCalleeAndAccessorKey(path.node.callee)
+
+					if (!callee) {
+						throw new Error('Callee is undefined. CODE: ' + generate(path.node).code)
 					}
 
 					const newNode = callExpression(identifier(useFunctionName), [
-						// @ts-ignore
-						transformCallee(path.node.callee),
+						callee,
 						objectExpression([
 							objectProperty(identifier('id'), stringLiteral(functionCallId)),
 							// @ts-ignore
 							objectProperty(identifier('args'), arrayExpression(path.node.arguments)),
-							objectProperty(identifier('name'), stringLiteral(functionCallName))
+							objectProperty(
+								identifier('name'),
+								// @ts-expect-error
+								accessorKey ? accessorKey : stringLiteral(functionCallName)
+							)
 						])
 					])
+
 					path.replaceWith(newNode)
 				}
 			})

test/testUtils.ts

@@ -1,2 +1,6 @@
 export const addAsync = (isAsync: boolean) => (isAsync ? 'async ' : '')
 export const addAwait = (isAsync: boolean) => (isAsync ? 'await ' : '')
+
+export function toOneLine(code: string) {
+	return code.split('\n').join('').replace(/\s+/g, '').replaceAll("'", '"').replaceAll(';', '')
+}