Add custom function signature parsing with flexible event extraction

Author: cursoragent

src/analyze/go/eventExtractor.js

@@ -145,8 +145,8 @@ function extractSnowplowEventName(callNode) {
  * @returns {string|null} Event name or null if not found
  */
 function extractCustomEventName(callNode) {
-  if (callNode.args && callNode.args.length > 0) {
-    return extractStringValue(callNode.args[0]);
+  if (callNode.args && callNode.args.length > 1) {
+    return extractStringValue(callNode.args[1]);
   }
   return null;
 }

src/analyze/go/index.js

@@ -8,16 +8,21 @@ const { extractGoAST } = require('./goAstParser');
 const { buildTypeContext } = require('./typeContext');
 const { deduplicateEvents } = require('./eventDeduplicator');
 const { extractEventsFromBody } = require('./astTraversal');
+const { processGoFile } = require('./utils');
+const { parseCustomFunctionSignature } = require('../utils/customFunctionParser');
 
 /**
  * Analyze a Go file and extract tracking events
  * @param {string} filePath - Path to the Go file to analyze
- * @param {string|null} customFunction - Name of custom tracking function to detect (optional)
+ * @param {string|null} customFunctionSignature - Signature of custom tracking function to detect (optional)
  * @returns {Promise<Array>} Array of tracking events found in the file
  * @throws {Error} If the file cannot be read or parsed
  */
-async function analyzeGoFile(filePath, customFunction) {
+async function analyzeGoFile(filePath, customFunctionSignature) {
   try {
+    const customConfig = customFunctionSignature ? parseCustomFunctionSignature(customFunctionSignature) : null;
+    const customFunction = customConfig ? customConfig.functionName : null;
+
     // Read the Go file
     const source = fs.readFileSync(filePath, 'utf8');
 

src/analyze/go/propertyExtractor.js

@@ -272,8 +272,24 @@ function extractSnowplowProperties(callNode, properties, typeContext, currentFun
  * @param {string} currentFunction - Current function context for type lookups
  */
 function extractCustomProperties(callNode, properties, typeContext, currentFunction) {
-  if (callNode.args && callNode.args.length > 1) {
-    extractPropertiesFromExpr(callNode.args[1], properties, typeContext, currentFunction);
+  if (!callNode.args || callNode.args.length === 0) return;
+
+  // Add userId from first argument if present
+  const userIdArg = callNode.args[0];
+  if (userIdArg) {
+    properties['userId'] = { type: 'string' };
+  }
+
+  // Properties map is expected to be third argument in new signature, fallback to second for legacy
+  let propsArg = null;
+  if (callNode.args.length >= 3) {
+    propsArg = callNode.args[2];
+  } else if (callNode.args.length >= 2) {
+    propsArg = callNode.args[1];
+  }
+
+  if (propsArg) {
+    extractPropertiesFromExpr(propsArg, properties, typeContext, currentFunction);
   }
 }
 

src/analyze/javascript/extractors/event-extractor.js

@@ -20,17 +20,22 @@ const EXTRACTION_STRATEGIES = {
   googleanalytics: extractGoogleAnalyticsEvent,
   snowplow: extractSnowplowEvent,
   mparticle: extractMparticleEvent,
+  custom: extractCustomEvent,
   default: extractDefaultEvent
 };
 
 /**
  * Extracts event information from a CallExpression node
  * @param {Object} node - AST CallExpression node
  * @param {string} source - Analytics provider source
+ * @param {Object} customConfig - Parsed custom function configuration
  * @returns {EventData} Extracted event data
  */
-function extractEventData(node, source) {
+function extractEventData(node, source, customConfig) {
   const strategy = EXTRACTION_STRATEGIES[source] || EXTRACTION_STRATEGIES.default;
+  if (source === 'custom') {
+    return strategy(node, customConfig);
+  }
   return strategy(node);
 }
 
@@ -113,16 +118,41 @@ function extractDefaultEvent(node) {
   return { eventName, propertiesNode };
 }
 
+/**
+ * Extracts Custom function event data according to signature
+ * @param {Object} node - CallExpression node
+ * @param {Object} customConfig - Parsed custom function configuration
+ * @returns {EventData & {extraArgs:Object}} event data plus extra args map
+ */
+function extractCustomEvent(node, customConfig) {
+  const args = node.arguments || [];
+
+  const eventArg = args[customConfig?.eventIndex ?? 0];
+  const propertiesArg = args[customConfig?.propertiesIndex ?? 1];
+
+  const eventName = getStringValue(eventArg);
+
+  const extraArgs = {};
+  if (customConfig && customConfig.extraParams) {
+    customConfig.extraParams.forEach(extra => {
+      extraArgs[extra.name] = args[extra.idx];
+    });
+  }
+
+  return { eventName, propertiesNode: propertiesArg, extraArgs };
+}
+
 /**
  * Processes extracted event data into final event object
  * @param {EventData} eventData - Raw event data
  * @param {string} source - Analytics source
  * @param {string} filePath - File path
  * @param {number} line - Line number
  * @param {string} functionName - Containing function name
+ * @param {Object} customConfig - Parsed custom function configuration
  * @returns {Object|null} Processed event object or null
  */
-function processEventData(eventData, source, filePath, line, functionName) {
+function processEventData(eventData, source, filePath, line, functionName, customConfig) {
   const { eventName, propertiesNode } = eventData;
 
   if (!eventName || !propertiesNode || propertiesNode.type !== NODE_TYPES.OBJECT_EXPRESSION) {
@@ -131,6 +161,15 @@ function processEventData(eventData, source, filePath, line, functionName) {
 
   let properties = extractProperties(propertiesNode);
 
+  // Handle custom extra params
+  if (source === 'custom' && customConfig && eventData.extraArgs) {
+    for (const [paramName, argNode] of Object.entries(eventData.extraArgs)) {
+      properties[paramName] = {
+        type: inferNodeValueType(argNode)
+      };
+    }
+  }
+
   // Special handling for Snowplow: remove 'action' from properties
   if (source === 'snowplow' && properties.action) {
     delete properties.action;
@@ -173,6 +212,25 @@ function findPropertyByKey(objectNode, key) {
   );
 }
 
+/**
+ * Infers the type of a value from an AST node (simple heuristic)
+ * @param {Object} node - AST node
+ * @returns {string} inferred type
+ */
+function inferNodeValueType(node) {
+  if (!node) return 'any';
+  switch (node.type) {
+    case NODE_TYPES.LITERAL:
+      return typeof node.value;
+    case NODE_TYPES.OBJECT_EXPRESSION:
+      return 'object';
+    case NODE_TYPES.ARRAY_EXPRESSION:
+      return 'array';
+    default:
+      return 'any';
+  }
+}
+
 module.exports = {
   extractEventData,
   processEventData

src/analyze/javascript/index.js

@@ -4,22 +4,24 @@
  */
 
 const { parseFile, findTrackingEvents, FileReadError, ParseError } = require('./parser');
+const { parseCustomFunctionSignature } = require('../utils/customFunctionParser');
 
 /**
  * Analyzes a JavaScript file for analytics tracking calls
  * @param {string} filePath - Path to the JavaScript file to analyze
  * @param {string} [customFunction] - Optional custom function name to detect
  * @returns {Array<Object>} Array of tracking events found in the file
  */
-function analyzeJsFile(filePath, customFunction) {
+function analyzeJsFile(filePath, customFunctionSignature) {
   const events = [];
+  const customConfig = customFunctionSignature ? parseCustomFunctionSignature(customFunctionSignature) : null;
 
   try {
     // Parse the file into an AST
     const ast = parseFile(filePath);
 
     // Find and extract tracking events
-    const foundEvents = findTrackingEvents(ast, filePath, customFunction);
+    const foundEvents = findTrackingEvents(ast, filePath, customConfig);
     events.push(...foundEvents);
 
   } catch (error) {

src/analyze/javascript/parser.js

@@ -70,16 +70,16 @@ function parseFile(filePath) {
  * Walks the AST and finds analytics tracking calls
  * @param {Object} ast - Parsed AST
  * @param {string} filePath - Path to the file being analyzed
- * @param {string} [customFunction] - Custom function name to detect
+ * @param {Object} [customConfig] - Custom function configuration object
  * @returns {Array<Object>} Array of found events
  */
-function findTrackingEvents(ast, filePath, customFunction) {
+function findTrackingEvents(ast, filePath, customConfig) {
   const events = [];
 
   walk.ancestor(ast, {
     [NODE_TYPES.CALL_EXPRESSION]: (node, ancestors) => {
       try {
-        const event = extractTrackingEvent(node, ancestors, filePath, customFunction);
+        const event = extractTrackingEvent(node, ancestors, filePath, customConfig);
         if (event) {
           events.push(event);
         }
@@ -97,25 +97,25 @@ function findTrackingEvents(ast, filePath, customFunction) {
  * @param {Object} node - CallExpression node
  * @param {Array<Object>} ancestors - Ancestor nodes
  * @param {string} filePath - File path
- * @param {string} [customFunction] - Custom function name
+ * @param {Object} [customConfig] - Custom function configuration object
  * @returns {Object|null} Extracted event or null
  */
-function extractTrackingEvent(node, ancestors, filePath, customFunction) {
+function extractTrackingEvent(node, ancestors, filePath, customConfig) {
   // Detect the analytics source
-  const source = detectAnalyticsSource(node, customFunction);
+  const source = detectAnalyticsSource(node, customConfig?.functionName);
   if (source === 'unknown') {
     return null;
   }
 
   // Extract event data based on the source
-  const eventData = extractEventData(node, source);
+  const eventData = extractEventData(node, source, customConfig);
 
   // Get location and context information
   const line = node.loc.start.line;
   const functionName = findWrappingFunction(node, ancestors);
 
   // Process the event data into final format
-  return processEventData(eventData, source, filePath, line, functionName);
+  return processEventData(eventData, source, filePath, line, functionName, customConfig);
 }
 
 module.exports = {

src/analyze/python/index.js

@@ -5,6 +5,7 @@
 
 const fs = require('fs');
 const path = require('path');
+const { parseCustomFunctionSignature } = require('../utils/customFunctionParser');
 
 // Singleton instance of Pyodide
 let pyodide = null;
@@ -40,7 +41,7 @@ async function initPyodide() {
  * libraries, extracting event names, properties, and metadata.
  * 
  * @param {string} filePath - Path to the Python file to analyze
- * @param {string} [customFunction=null] - Name of a custom tracking function to detect
+ * @param {string} [customFunctionSignature=null] - Signature of a custom tracking function to detect
  * @returns {Promise<Array<Object>>} Array of tracking events found in the file
  * @returns {Promise<Array>} Empty array if an error occurs
  * 
@@ -52,7 +53,10 @@ async function initPyodide() {
  * // With custom tracking function
  * const events = await analyzePythonFile('./app.py', 'track_event');
  */
-async function analyzePythonFile(filePath, customFunction = null) {
+async function analyzePythonFile(filePath, customFunctionSignature = null) {
+  const customConfig = customFunctionSignature ? parseCustomFunctionSignature(customFunctionSignature) : null;
+  const customFunction = customConfig ? customConfig.functionName : null;
+
   // Validate inputs
   if (!filePath || typeof filePath !== 'string') {
     console.error('Invalid file path provided');

src/analyze/python/pythonTrackingAnalyzer.py

@@ -430,7 +430,10 @@ def _extract_snowplow_event_name(self, node: ast.Call) -> Optional[str]:
 
     def _extract_custom_event_name(self, node: ast.Call) -> Optional[str]:
         """Extract event name for custom tracking function."""
-        # Standard format: customFunction('event_name', {...})
+        # Support signatures where event name might be second argument (userId first)
+        if len(node.args) >= 2 and isinstance(node.args[1], ast.Constant):
+            return node.args[1].value
+        # Fallback to legacy format: first argument is event name
         if len(node.args) >= 1 and isinstance(node.args[0], ast.Constant):
             return node.args[0].value
         return None
@@ -502,6 +505,10 @@ def _extract_user_id(self, node: ast.Call, source: str) -> EventProperties:
             # Check if event is not anonymous and extract distinct_id
             user_id_props.update(self._extract_posthog_user_id(node))
 
+        elif source == 'custom':
+            if len(node.args) >= 1 and self._is_non_null_value(node.args[0]):
+                user_id_props["userId"] = {"type": "string"}
+            
         return user_id_props
 
     def _is_non_null_value(self, node: ast.AST) -> bool:
@@ -559,10 +566,13 @@ def _get_properties_node(self, node: ast.Call, source: str) -> Optional[ast.Dict
                         return keyword.value
 
         elif source == 'custom':
-            # Properties are in the second argument
-            if len(node.args) > 1:
+            # New custom function format: customFunction(userId, 'event_name', {...})
+            if len(node.args) >= 3 and isinstance(node.args[2], ast.Dict):
+                return node.args[2]
+            # Fallback legacy two-arg format
+            if len(node.args) >= 2 and isinstance(node.args[1], ast.Dict):
                 return node.args[1]
-                
+            
         elif source == 'posthog':
             # Check named parameters first, then positional
             for keyword in node.keywords: