Custom Function Signatures

README.md

@@ -23,38 +23,39 @@ Run without installation! Just use:
 npx @flisk/analyze-tracking /path/to/project [options]
 ```
 
-### Key Options:
+### Key Options
 - `-g, --generateDescription`: Generate descriptions of fields (default: `false`)
 - `-p, --provider <provider>`: Specify a provider (options: `openai`, `gemini`)
 - `-m, --model <model>`: Specify a model (ex: `gpt-4.1-nano`, `gpt-4o-mini`, `gemini-2.0-flash-lite-001`)
 - `-o, --output <output_file>`: Name of the output file (default: `tracking-schema.yaml`)
-- `-c, --customFunction <function_name>`: Specify a custom tracking function
+- `-c, --customFunction <function_signature>`: Specify the signature of your custom tracking function (see [instructions here](#custom-functions))
 - `--format <format>`: Output format, either `yaml` (default) or `json`. If an invalid value is provided, the CLI will exit with an error.
 - `--stdout`: Print the output to the terminal instead of writing to a file (works with both YAML and JSON)
 
 🔑&nbsp; **Important:** If you are using `generateDescription`, you must set the appropriate credentials for the LLM provider you are using as an environment variable. OpenAI uses `OPENAI_API_KEY` and Google Vertex AI uses `GOOGLE_APPLICATION_CREDENTIALS`.
 
-<details>
-  <summary>Note on Custom Functions 💡</summary>
 
-  Use this if you have your own in-house tracker or a wrapper function that calls other tracking libraries.
+### Custom Functions
 
-  We currently only support functions that follow the following format:
-
-  **JavaScript/TypeScript/Python/Ruby:**
-  ```js
-  yourCustomTrackFunctionName('<event_name>', {
-    <event_parameters>
-  });
-  ```
-
-  **Go:**
-  ```go
-  yourCustomTrackFunctionName("<event_name>", map[string]any{}{
-    "<property_name>": "<property_value>",
-  })
-  ```
-</details>
+If you have your own in-house tracker or a wrapper function that calls other tracking libraries, you can specify the function signature with the `-c` or `--customFunction` option.
+
+Your function signature should be in the following format:
+```js
+yourCustomTrackFunctionName(EVENT_NAME, PROPERTIES, customFieldOne, customFieldTwo)
+```
+
+- `EVENT_NAME` is the name of the event you are tracking. It should be a string or a pointer to a string. This is required.
+- `PROPERTIES` is an object of properties for that event. It should be an object / dictionary. This is optional.
+- Any additional parameters are other fields you are tracking. They can be of any type. The names you provide for these parameters will be used as the property names in the output.
+
+
+For example, if your function has a userId parameter at the beginning, followed by the event name and properties, you would pass in the following:
+
+```js
+yourCustomTrackFunctionName(userId, EVENT_NAME, PROPERTIES)
+```
+
+If your function follows the standard format `yourCustomTrackFunctionName(EVENT_NAME, PROPERTIES)`, you can simply pass in `yourCustomTrackFunctionName` to `--customFunction` as a shorthand.
 
 
 ## What's Generated?

src/analyze/go/astTraversal.js

@@ -12,34 +12,34 @@ const { extractTrackingEvent } = require('./trackingExtractor');
  * @param {Array<Object>} events - Array to collect found tracking events (modified in place)
  * @param {string} filePath - Path to the file being analyzed
  * @param {string} functionName - Name of the current function being processed
- * @param {string|null} customFunction - Name of custom tracking function to detect
+ * @param {Object|null} customConfig - Parsed custom function configuration (or null)
  * @param {Object} typeContext - Type information context for variable resolution
  * @param {string} currentFunction - Current function context for type lookups
  */
-function extractEventsFromBody(body, events, filePath, functionName, customFunction, typeContext, currentFunction) {
+function extractEventsFromBody(body, events, filePath, functionName, customConfig, typeContext, currentFunction) {
   for (const stmt of body) {
     if (stmt.tag === 'exec' && stmt.expr) {
-      processExpression(stmt.expr, events, filePath, functionName, customFunction, typeContext, currentFunction);
+      processExpression(stmt.expr, events, filePath, functionName, customConfig, typeContext, currentFunction);
     } else if (stmt.tag === 'declare' && stmt.value) {
       // Handle variable declarations with tracking calls
-      processExpression(stmt.value, events, filePath, functionName, customFunction, typeContext, currentFunction);
+      processExpression(stmt.value, events, filePath, functionName, customConfig, typeContext, currentFunction);
     } else if (stmt.tag === 'assign' && stmt.rhs) {
       // Handle assignments with tracking calls
-      processExpression(stmt.rhs, events, filePath, functionName, customFunction, typeContext, currentFunction);
+      processExpression(stmt.rhs, events, filePath, functionName, customConfig, typeContext, currentFunction);
     } else if (stmt.tag === 'if' && stmt.body) {
-      extractEventsFromBody(stmt.body, events, filePath, functionName, customFunction, typeContext, currentFunction);
+      extractEventsFromBody(stmt.body, events, filePath, functionName, customConfig, typeContext, currentFunction);
     } else if (stmt.tag === 'elseif' && stmt.body) {
-      extractEventsFromBody(stmt.body, events, filePath, functionName, customFunction, typeContext, currentFunction);
+      extractEventsFromBody(stmt.body, events, filePath, functionName, customConfig, typeContext, currentFunction);
     } else if (stmt.tag === 'else' && stmt.body) {
-      extractEventsFromBody(stmt.body, events, filePath, functionName, customFunction, typeContext, currentFunction);
+      extractEventsFromBody(stmt.body, events, filePath, functionName, customConfig, typeContext, currentFunction);
     } else if (stmt.tag === 'for' && stmt.body) {
-      extractEventsFromBody(stmt.body, events, filePath, functionName, customFunction, typeContext, currentFunction);
+      extractEventsFromBody(stmt.body, events, filePath, functionName, customConfig, typeContext, currentFunction);
     } else if (stmt.tag === 'foreach' && stmt.body) {
-      extractEventsFromBody(stmt.body, events, filePath, functionName, customFunction, typeContext, currentFunction);
+      extractEventsFromBody(stmt.body, events, filePath, functionName, customConfig, typeContext, currentFunction);
     } else if (stmt.tag === 'switch' && stmt.cases) {
       for (const caseNode of stmt.cases) {
         if (caseNode.body) {
-          extractEventsFromBody(caseNode.body, events, filePath, functionName, customFunction, typeContext, currentFunction);
+          extractEventsFromBody(caseNode.body, events, filePath, functionName, customConfig, typeContext, currentFunction);
         }
       }
     }
@@ -52,44 +52,44 @@ function extractEventsFromBody(body, events, filePath, functionName, customFunct
  * @param {Array<Object>} events - Array to collect found tracking events (modified in place)
  * @param {string} filePath - Path to the file being analyzed
  * @param {string} functionName - Name of the current function being processed
- * @param {string|null} customFunction - Name of custom tracking function to detect
+ * @param {Object|null} customConfig - Parsed custom function configuration (or null)
  * @param {Object} typeContext - Type information context for variable resolution
  * @param {string} currentFunction - Current function context for type lookups
  * @param {number} [depth=0] - Current recursion depth (used to prevent infinite recursion)
  */
-function processExpression(expr, events, filePath, functionName, customFunction, typeContext, currentFunction, depth = 0) {
+function processExpression(expr, events, filePath, functionName, customConfig, typeContext, currentFunction, depth = 0) {
   if (!expr || depth > MAX_RECURSION_DEPTH) return; // Prevent infinite recursion with depth limit
 
   // Handle array of expressions
   if (Array.isArray(expr)) {
     for (const item of expr) {
-      processExpression(item, events, filePath, functionName, customFunction, typeContext, currentFunction, depth + 1);
+      processExpression(item, events, filePath, functionName, customConfig, typeContext, currentFunction, depth + 1);
     }
     return;
   }
 
   // Handle single expression with body
   if (expr.body) {
     for (const item of expr.body) {
-      processExpression(item, events, filePath, functionName, customFunction, typeContext, currentFunction, depth + 1);
+      processExpression(item, events, filePath, functionName, customConfig, typeContext, currentFunction, depth + 1);
     }
     return;
   }
 
   // Handle specific node types
   if (expr.tag === 'call') {
-    const trackingCall = extractTrackingEvent(expr, filePath, functionName, customFunction, typeContext, currentFunction);
+    const trackingCall = extractTrackingEvent(expr, filePath, functionName, customConfig, typeContext, currentFunction);
     if (trackingCall) {
       events.push(trackingCall);
     }
 
     // Also process call arguments
     if (expr.args) {
-      processExpression(expr.args, events, filePath, functionName, customFunction, typeContext, currentFunction, depth + 1);
+      processExpression(expr.args, events, filePath, functionName, customConfig, typeContext, currentFunction, depth + 1);
     }
   } else if (expr.tag === 'structlit') {
     // Check if this struct literal is a tracking event
-    const trackingCall = extractTrackingEvent(expr, filePath, functionName, customFunction, typeContext, currentFunction);
+    const trackingCall = extractTrackingEvent(expr, filePath, functionName, customConfig, typeContext, currentFunction);
     if (trackingCall) {
       events.push(trackingCall);
     }
@@ -98,21 +98,21 @@ function processExpression(expr, events, filePath, functionName, customFunction,
     if (!trackingCall && expr.fields) {
       for (const field of expr.fields) {
         if (field.value) {
-          processExpression(field.value, events, filePath, functionName, customFunction, typeContext, currentFunction, depth + 1);
+          processExpression(field.value, events, filePath, functionName, customConfig, typeContext, currentFunction, depth + 1);
         }
       }
     }
   }
 
   // Process other common properties that might contain expressions
   if (expr.value && expr.tag !== 'structlit') {
-    processExpression(expr.value, events, filePath, functionName, customFunction, typeContext, currentFunction, depth + 1);
+    processExpression(expr.value, events, filePath, functionName, customConfig, typeContext, currentFunction, depth + 1);
   }
   if (expr.lhs) {
-    processExpression(expr.lhs, events, filePath, functionName, customFunction, typeContext, currentFunction, depth + 1);
+    processExpression(expr.lhs, events, filePath, functionName, customConfig, typeContext, currentFunction, depth + 1);
   }
   if (expr.rhs) {
-    processExpression(expr.rhs, events, filePath, functionName, customFunction, typeContext, currentFunction, depth + 1);
+    processExpression(expr.rhs, events, filePath, functionName, customConfig, typeContext, currentFunction, depth + 1);
   }
 }
 

src/analyze/go/eventExtractor.js

@@ -10,9 +10,10 @@ const { extractStringValue, findStructLiteral, findStructField, extractSnowplowV
  * Extract event name from a tracking call based on the source
  * @param {Object} callNode - AST node representing a function call or struct literal
  * @param {string} source - Analytics source (e.g., 'segment', 'amplitude')
+ * @param {Object|null} customConfig - Parsed custom function configuration
  * @returns {string|null} Event name or null if not found
  */
-function extractEventName(callNode, source) {
+function extractEventName(callNode, source, customConfig = null) {
   if (!callNode.args || callNode.args.length === 0) {
     // For struct literals, we need to check fields instead of args
     if (!callNode.fields || callNode.fields.length === 0) {
@@ -35,7 +36,7 @@ function extractEventName(callNode, source) {
       return extractSnowplowEventName(callNode);
 
     case ANALYTICS_SOURCES.CUSTOM:
-      return extractCustomEventName(callNode);
+      return extractCustomEventName(callNode, customConfig);
   }
 
   return null;
@@ -142,13 +143,15 @@ function extractSnowplowEventName(callNode) {
  * Extract custom event name
  * Pattern: customFunction("event_name", props)
  * @param {Object} callNode - AST node for custom tracking function call
+ * @param {Object|null} customConfig - Custom configuration object
  * @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]);
-  }
-  return null;
+function extractCustomEventName(callNode, customConfig) {
+  if (!callNode.args || callNode.args.length === 0) return null;
+  const args = callNode.args;
+  const eventIdx = customConfig?.eventIndex ?? 0;
+  const argNode = args[eventIdx];
+  return extractStringValue(argNode);
 }
 
 module.exports = {

src/analyze/go/index.js

@@ -8,16 +8,20 @@ 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;
+
     // Read the Go file
     const source = fs.readFileSync(filePath, 'utf8');
 
@@ -37,7 +41,7 @@ async function analyzeGoFile(filePath, customFunction) {
         currentFunction = node.name;
         // Process the function body
         if (node.body) {
-          extractEventsFromBody(node.body, events, filePath, currentFunction, customFunction, typeContext, currentFunction);
+          extractEventsFromBody(node.body, events, filePath, currentFunction, customConfig, typeContext, currentFunction);
         }
       }
     }

src/analyze/go/propertyExtractor.js

@@ -12,9 +12,10 @@ const { extractStringValue, findStructLiteral, findStructField, extractFieldName
  * @param {string} source - Analytics source (e.g., 'segment', 'amplitude')
  * @param {Object} typeContext - Type information context for variable resolution
  * @param {string} currentFunction - Current function context for type lookups
+ * @param {Object} customConfig - Custom configuration for property extraction
  * @returns {Object} Object containing extracted properties with their type information
  */
-function extractProperties(callNode, source, typeContext, currentFunction) {
+function extractProperties(callNode, source, typeContext, currentFunction, customConfig) {
   const properties = {};
 
   switch (source) {
@@ -36,7 +37,7 @@ function extractProperties(callNode, source, typeContext, currentFunction) {
       break;
 
     case ANALYTICS_SOURCES.CUSTOM:
-      extractCustomProperties(callNode, properties, typeContext, currentFunction);
+      extractCustomProperties(callNode, properties, typeContext, currentFunction, customConfig);
       break;
   }
 
@@ -270,10 +271,29 @@ function extractSnowplowProperties(callNode, properties, typeContext, currentFun
  * @param {Object} properties - Object to store extracted properties (modified in place)
  * @param {Object} typeContext - Type information context for variable resolution
  * @param {string} currentFunction - Current function context for type lookups
+ * @param {Object} customConfig - Custom configuration for property extraction
  */
-function extractCustomProperties(callNode, properties, typeContext, currentFunction) {
-  if (callNode.args && callNode.args.length > 1) {
-    extractPropertiesFromExpr(callNode.args[1], properties, typeContext, currentFunction);
+function extractCustomProperties(callNode, properties, typeContext, currentFunction, customConfig) {
+  if (!callNode.args || callNode.args.length === 0) return;
+
+  const args = callNode.args;
+
+  const propsIdx = customConfig?.propertiesIndex ?? 1;
+
+  // Extract extra params first (those not event or properties)
+  if (customConfig && Array.isArray(customConfig.extraParams)) {
+    customConfig.extraParams.forEach(param => {
+      const argNode = args[param.idx];
+      if (argNode) {
+        properties[param.name] = getPropertyInfo(argNode, typeContext, currentFunction);
+      }
+    });
+  }
+
+  // Extract properties map/object (if provided)
+  const propsArg = args[propsIdx];
+  if (propsArg) {
+    extractPropertiesFromExpr(propsArg, properties, typeContext, currentFunction);
   }
 }
 

src/analyze/go/trackingExtractor.js

@@ -13,19 +13,19 @@ const { extractProperties } = require('./propertyExtractor');
  * @param {Object} callNode - AST node representing a function call or struct literal
  * @param {string} filePath - Path to the file being analyzed
  * @param {string} functionName - Name of the function containing this tracking call
- * @param {string|null} customFunction - Name of custom tracking function to detect
+ * @param {Object|null} customConfig - Parsed custom function configuration (or null)
  * @param {Object} typeContext - Type information context for variable resolution
  * @param {string} currentFunction - Current function context for type lookups
  * @returns {Object|null} Tracking event object with eventName, source, properties, etc., or null if not a tracking call
  */
-function extractTrackingEvent(callNode, filePath, functionName, customFunction, typeContext, currentFunction) {
-  const source = detectSource(callNode, customFunction);
+function extractTrackingEvent(callNode, filePath, functionName, customConfig, typeContext, currentFunction) {
+  const source = detectSource(callNode, customConfig ? customConfig.functionName : null);
   if (!source) return null;
 
-  const eventName = extractEventName(callNode, source);
+  const eventName = extractEventName(callNode, source, customConfig);
   if (!eventName) return null;
 
-  const properties = extractProperties(callNode, source, typeContext, currentFunction);
+  const properties = extractProperties(callNode, source, typeContext, currentFunction, customConfig);
 
   // Get line number based on source type
   let line = 0;