feat: support custom NodejsFunction constructs

README.md

@@ -4,7 +4,8 @@ A library that delivers tools to analyze the bundle size of TypeScript/JavaScrip
 
 ## Prerequisites 📓
 
-1. Use the [NodejsFunction construct](https://docs.aws.amazon.com/cdk/api/v2/docs/aws-cdk-lib.aws_lambda_nodejs-readme.html) (or a custom construct that ends up using it) to define functions
+1. Use the [`NodejsFunction` construct](https://docs.aws.amazon.com/cdk/api/v2/docs/aws-cdk-lib.aws_lambda_nodejs-readme.html) to define functions
+   - If using a custom construct that extends the `NodejsFunction` construct, refer to the [Custom `NodejsFunction` construct section](#custom-nodejsfunction-construct-🏗️)
 2. Have [local bundling](https://docs.aws.amazon.com/cdk/api/v2/docs/aws-cdk-lib.aws_lambda_nodejs-readme.html#local-bundling) enabled. Basically, this means that the bundling of the lambdas is done on the machine, not inside a docker container. To achieve that, `esbuild` must be installed in the project. Follow the link above for more details.
 
 ## Usage 📦
@@ -70,6 +71,32 @@ Example:
 cdk synth --quiet -c analyze=Health -c template=sunburst
 ```
 
+## Custom `NodejsFunction` construct 🏗️
+
+If using a custom construct that extends the `NodejsFunction` construct, simply pass the custom construct that extends the `NodejsFunction` construct to the `NodeJsFunctionBundleAnalyzerAspect` constructor.
+
+For example:
+
+```ts
+import { NodeJsFunctionBundleAnalyzerAspect } from 'cdk-bundle-analyzer';
+
+const app = new App();
+
+// ...
+
+class MyCustomNodejsFunction extends NodejsFunction {
+  // ...
+}
+
+// ...
+
+Aspects.of(app).add(
+  new NodeJsFunctionBundleAnalyzerAspect({
+    customFunctionConstruct: MyCustomNodejsFunction,
+  }),
+);
+```
+
 ## Remarks 📝
 
 - The `NodeJsFunctionBundleAnalyzerAspect` will have no effect on the CDK app whatsoever. It will not change the behavior of the CDK app in any way. Moreover, the side-effect that generates the bundle analysis will only be executed if the `analyze` context variable is specified. Thus it is safe to add the aspect to the CDK app and commit it to the repository.

src/NodeJsFunctionBundleAnalyzerAspect.ts

@@ -9,17 +9,33 @@ import { tmpdir } from 'os';
 import { dirname, join } from 'path';
 
 /**
- * Return `true` if `node` is a `NodejsFunction` instance.
+ * Return `true` if `node` is a `NodejsFunction` instance or a custom function construct instance.
  *
  * We cannot use `instanceof` because the `NodejsFunction` class will be copied in the imported code.
  *
  * See https://stackoverflow.com/a/63937850
  */
-function isNodejsFunction(node: IConstruct): node is NodejsFunction {
-  return node.constructor.name === NodejsFunction.name;
+function isNodejsFunction(
+  node: IConstruct,
+  customFunctionConstructName?: string,
+): node is NodejsFunction {
+  return (
+    node.constructor.name === NodejsFunction.name ||
+    node.constructor.name === customFunctionConstructName
+  );
+}
+
+interface NodeJsFunctionBundleAnalyzerAspectProps {
+  customFunctionConstruct?: typeof NodejsFunction;
 }
 
 class NodeJsFunctionBundleAnalyzerAspect implements IAspect {
+  private customFunctionConstructName: string | undefined;
+
+  constructor({ customFunctionConstruct }: NodeJsFunctionBundleAnalyzerAspectProps) {
+    this.customFunctionConstructName = customFunctionConstruct?.prototype.constructor.name;
+  }
+
   async visit(node: IConstruct): Promise<void> {
     const functionToAnalyze = node.node.tryGetContext('analyze') as string | undefined;
     if (functionToAnalyze === undefined) {
@@ -33,18 +49,12 @@ class NodeJsFunctionBundleAnalyzerAspect implements IAspect {
       );
     }
 
-    if (isNodejsFunction(node)) {
-      const functionName = node
-        .toString()
-        .replace(node.stack.stackName, '')
-        .replace('Lambda', '')
-        .replace(/\//g, '');
-
-      if (functionName !== functionToAnalyze) {
+    if (isNodejsFunction(node, this.customFunctionConstructName)) {
+      if (!node.toString().includes(functionToAnalyze)) {
         return;
       }
 
-      console.log(`\n⏳ Analyzing function ${functionName}`);
+      console.log(`\n⏳ Analyzing function ${functionToAnalyze}`);
 
       const assetPath = (node.node.defaultChild as CfnFunction).cfnOptions.metadata?.[
         'aws:asset:path'
@@ -58,7 +68,7 @@ class NodeJsFunctionBundleAnalyzerAspect implements IAspect {
         });
       } catch (e) {
         console.error(
-          `\n🤯 Analyze failed: metafile ${metafilePath} not found. Did you set metafile: true in the bundling options of the ${functionName} NodejsFunction construct?\n`,
+          `\n🤯 Analyze failed: metafile ${metafilePath} not found. Did you set metafile: true in the bundling options of the ${functionToAnalyze} NodejsFunction construct?\n`,
         );
 
         throw new Error('Analyze failed');
@@ -67,14 +77,14 @@ class NodeJsFunctionBundleAnalyzerAspect implements IAspect {
       const jsonContent = JSON.parse(textContent) as Metadata;
 
       const fileContent = await visualizer(jsonContent, {
-        title: `${functionName} function bundle visualizer `,
+        title: `${functionToAnalyze} function bundle visualizer `,
         template: template ?? 'treemap',
       });
 
       const TMP_FOLDER = join(tmpdir(), 'cdk-bundle-analyzer');
       const TEMP_DIR_LOCATION = join(TMP_FOLDER, new Date().getTime().toString());
 
-      const filename = `${TEMP_DIR_LOCATION}/${functionName}.html`;
+      const filename = `${TEMP_DIR_LOCATION}/${functionToAnalyze}.html`;
 
       await fs.mkdir(dirname(filename), { recursive: true });
       await fs.writeFile(filename, fileContent);