chore: set PYTHONPATH for workspace modules

Author: eoinsha

src/bundling.ts

@@ -23,26 +23,6 @@ export const DEFAULT_ASSET_EXCLUDES = [
   '.git/',
 ];
 
-const userCustomizePy = `
-import sys
-import os
-from pathlib import Path
-
-base_path = Path(__file__).parent
-
-for pth_file in base_path.glob('*.pth'):
-    print('Adding', pth_file, base_path)
-    with open(pth_file, 'r') as f:
-        for line in f:
-            module_path_str = str(base_path / line.strip())
-            print('...', module_path_str)
-            if module_path_str not in sys.path:
-                sys.path.insert(0, module_path_str)
-
-print('sys.path', sys.path)
-    
-`;
-
 interface BundlingCommandOptions {
   readonly rootDir: string;
   readonly workspacePackage?: string;
@@ -164,21 +144,16 @@ export class Bundling {
       ...options.commandHooks?.beforeBundling(options.inputDir, options.outputDir) ?? [],
     );
     commands.push(...[
-      `while [ -e ${options.inputDir}/wait.txt ]; do sleep 2; echo Waiting; done`,
       `rsync -rLv ${excludeArgs.join(' ')} ${options.inputDir}/ ${options.outputDir}`,
       `cd ${options.outputDir}`, // uv pip install needs to be run from here for editable deps to relative paths to be resolved
       `VIRTUAL_ENV=/tmp/venv uv sync ${uvCommonArgs} ${uvPackageArgs} --compile-bytecode --no-dev --frozen --no-editable --link-mode=copy`,
       `VIRTUAL_ENV=/tmp/venv uv export ${uvCommonArgs} ${uvPackageArgs} --no-dev --frozen --no-editable > ${reqsFile}`,
       `uv pip install -r ${reqsFile} --target ${options.outputDir} --reinstall --compile-bytecode --link-mode=copy --editable $(grep -e "^\./" ${reqsFile})`,
-      `sed -i 's|${options.outputDir}/|.|g' ${options.outputDir}/*.pth`,
       `rm -rf ${options.outputDir}/.venv`,
-      `echo ${Buffer.from(userCustomizePy).toString('base64')} | base64 -d > ${options.outputDir}/usercustomize.py`,
-      `while [ -e ${options.inputDir}/wait2.txt ]; do sleep 2; echo Waiting; done`,
     ]);
     commands.push(
       ...options.commandHooks?.afterBundling(options.inputDir, options.outputDir) ?? [],
     );
-    console.log('Bundling commands', { options, commands });
 
     return commands;
   }

src/function.ts

@@ -1,7 +1,9 @@
+import * as fs from 'node:fs';
 import * as path from 'node:path';
 import { Stack } from 'aws-cdk-lib';
 import {
   Architecture,
+  type CfnFunction,
   // biome-ignore lint/suspicious/noShadowRestrictedNames: shadows 'function'
   Function,
   type FunctionOptions,
@@ -57,16 +59,14 @@ export class PythonFunction extends Function {
     const {
       workspacePackage,
       handler = 'handler',
+      index = 'index.py',
       runtime = Runtime.PYTHON_3_12,
     } = props;
 
     const architecture = props.architecture ?? Architecture.ARM_64;
     const rootDir = path.resolve(props.rootDir);
 
-    let resolvedHandler = handler;
-    if (workspacePackage) {
-      resolvedHandler = `${workspacePackage.replace(/-/g, '_')}.${handler}`;
-    }
+    const resolvedHandler = `${index.slice(0, -3)}.${handler}`.replace(/\//g, '.');
 
     if (runtime.family !== RuntimeFamily.PYTHON) {
       throw new Error('Only Python runtimes are supported');
@@ -80,12 +80,40 @@ export class PythonFunction extends Function {
       workspacePackage,
       ...props.bundling,
     });
+
+    const environment = props.environment ?? {};
+
     super(scope, id, {
       ...props,
+      environment,
       architecture,
       runtime,
       code,
       handler: resolvedHandler,
     });
+
+    const assetPath = ((this.node.defaultChild) as CfnFunction).getMetadata('aws:asset:path');
+    const codePath = path.join(process.env.CDK_OUTDIR as string, assetPath);
+
+    const pythonPaths = getPthFilePaths(codePath);
+
+    if (pythonPaths.length > 0) {
+      let pythonPathValue = environment.PYTHONPATH;
+      const addedPaths = pythonPaths.join(':');
+      pythonPathValue = pythonPathValue ? `${pythonPathValue}:${addedPaths}` : addedPaths;
+      this.addEnvironment('PYTHONPATH', pythonPathValue);
+    }
+  }
+}
+
+function getPthFilePaths(basePath: string): string[] {
+  const pthFiles = fs.readdirSync(basePath).filter(file => file.endsWith('.pth'));
+  const pythonPaths: string[] = [];
+  for (const pthFile of pthFiles) {
+    const filePath = path.join(basePath, pthFile);
+    const content = fs.readFileSync(filePath, 'utf-8');
+    const dirs = content.split('\n').filter(line => line.trim() !== '');
+    pythonPaths.push(...dirs.map(dir => path.join('/var/task', path.relative('/asset-output', dir))));
   }
+  return pythonPaths;
 }