feat: handle package deps and pth

Author: eoinsha

.projenrc.ts

@@ -20,4 +20,5 @@ const project = new awscdk.AwsCdkConstructLibrary({
   devDeps: ['@biomejs/biome'] /* Build dependencies for this module. */,
   // packageName: undefined,  /* The "name" in package.json. */
 });
+project.files;
 project.synth();

API.md

@@ -2,23 +2,20 @@
 # The correct AWS SAM build image based on the runtime of the function will be
 # passed as build arg. The default allows to do `docker build .` when testing.
 ARG PYTHON_VERSION=3.7
-ARG UV_VERSION=0.4.20
 ARG IMAGE=public.ecr.aws/sam/build-python${PYTHON_VERSION}
 FROM $IMAGE
 
 ARG PIP_INDEX_URL
 ARG PIP_EXTRA_INDEX_URL
 ARG HTTPS_PROXY
-ARG POETRY_VERSION=1.5.1
+ARG UV_VERSION=0.4.20
 
 ENV PIP_CACHE_DIR=/tmp/pip-cache
 ENV UV_CACHE_DIR=/tmp/uv-cache
 
-# RUN mkdir /tmp/pip-cache && \
-#     chmod -R 777 /tmp/pip-cache && \
-#     mkdir /tmp/uv-cache && \
-#     chmod -R 777 /tmp/uv-cache && \
-#     pip install uv==${UV_VERSION} \
-#     rm -rf /tmp/uv-cache/* /tmp/pip-cache/*
+RUN mkdir /tmp/pip-cache && \
+    chmod -R 777 /tmp/pip-cache && \
+    pip install uv==$UV_VERSION && \
+    rm -rf /tmp/pip-cache/*
 
 CMD [ "python" ]

src/Dockerfile resources/Dockerfilesrc/Dockerfile renamed to resources/Dockerfile

@@ -1,9 +1,9 @@
+import * as path from 'node:path';
 import {
   AssetHashType,
   AssetStaging,
   type BundlingFileAccess,
   DockerImage,
-  type DockerRunOptions,
   type DockerVolume,
 } from 'aws-cdk-lib';
 import {
@@ -12,63 +12,83 @@ import {
   Code,
   type Runtime,
 } from 'aws-cdk-lib/aws-lambda';
+import type { BundlingOptions, ICommandHooks } from './types';
 
 export const HASHABLE_DEPENDENCIES_EXCLUDE = ['*.pyc'];
 
+export const DEFAULT_ASSET_EXCLUDES = [
+  '.venv/',
+  'node_modules/',
+  'cdk.out/',
+  '.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 entry: string;
+  readonly rootDir: string;
+  readonly workspacePackage?: string;
   readonly inputDir: string;
   readonly outputDir: string;
+  readonly assetExcludes: string[];
+  readonly commandHooks?: ICommandHooks;
 }
 
-export interface BundlingProps extends DockerRunOptions {
-  readonly entry: string;
-  readonly runtime: Runtime;
-
+export interface BundlingProps extends BundlingOptions {
   /**
-   * Lambda CPU architecture
-   *
-   * @default Architecture.X86_64
+   * uv project root (workspace root)
    */
-  readonly architecture?: Architecture;
+  readonly rootDir: string;
 
   /**
-   * Skip bundling process
-   *
-   * @default false
+   * uv package to use for the Lambda Function
    */
-  readonly skip?: boolean;
+  readonly workspacePackage?: string;
 
   /**
-   * Docker image to use for bundling.
-   *
-   * @default - Default bundling image from the sam/build-python set
+   * Lambda runtime (must be one of the Python runtimes)
    */
-  readonly image?: DockerImage;
+  readonly runtime: Runtime;
 
   /**
-   * Optional build arguments to pass to the bundling container when the default
-   * image is used.
+   * Lambda CPU architecture
    *
-   * @default - {}
+   * @default Architecture.ARM_64
    */
-  readonly buildArgs?: { [key: string]: string };
+  readonly architecture?: Architecture;
 
   /**
-   * Specifies how to copy files to/from the docker container. BIND_MOUNT is generally faster
-   * but VOLUME_MOUNT works with remote Docker contexts.
+   * Skip bundling process
    *
-   * @default - BundlingFileAccess.BIND_MOUNT
+   * @default false
    */
-  readonly bundlingFileAccess?: BundlingFileAccess;
+  readonly skip?: boolean;
 }
 
 /**
  * Bundling options for Python Lambda assets
  */
 export class Bundling {
   public static bundle(options: BundlingProps): AssetCode {
-    return Code.fromAsset(options.entry, {
+    return Code.fromAsset(options.rootDir, {
       assetHashType: AssetHashType.SOURCE,
       exclude: HASHABLE_DEPENDENCIES_EXCLUDE,
       bundling: options.skip ? undefined : new Bundling(options),
@@ -88,17 +108,28 @@ export class Bundling {
   public readonly bundlingFileAccess?: BundlingFileAccess | undefined;
 
   constructor(props: BundlingProps) {
-    const { entry, image, runtime, architecture = Architecture.X86_64 } = props;
+    const {
+      rootDir,
+      workspacePackage,
+      image,
+      runtime,
+      commandHooks,
+      assetExcludes = DEFAULT_ASSET_EXCLUDES,
+      architecture = Architecture.ARM_64,
+    } = props;
 
     const bundlingCommands = this.createBundlingCommands({
-      entry,
+      rootDir,
+      workspacePackage,
+      assetExcludes,
+      commandHooks,
       inputDir: AssetStaging.BUNDLING_INPUT_DIR,
       outputDir: AssetStaging.BUNDLING_OUTPUT_DIR,
     });
 
     this.image =
       image ??
-      DockerImage.fromBuild(__dirname, {
+      DockerImage.fromBuild(path.resolve(__dirname, '..', 'resources'), {
         buildArgs: {
           ...props.buildArgs,
           IMAGE: runtime.bundlingImage.image,
@@ -123,6 +154,32 @@ export class Bundling {
   }
 
   private createBundlingCommands(options: BundlingCommandOptions): string[] {
-    return [`rsync -rLv ${options.inputDir}/ ${options.outputDir}`, 'uv sync'];
+    const excludeArgs = options.assetExcludes.map((exclude) => `--exclude="${exclude}"`);
+    const workspacePackage = options.workspacePackage;
+    const uvCommonArgs = `--directory ${options.outputDir}`;
+    const uvPackageArgs = workspacePackage ? `--package ${workspacePackage}` : '';
+    const reqsFile = `/tmp/requirements${workspacePackage || ''}.txt`;
+    const commands = [];
+    commands.push(
+      ...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/bundling.ts

@@ -1,8 +1,8 @@
-import * as fs from 'node:fs';
 import * as path from 'node:path';
-import { type BundlingOptions, Stack } from 'aws-cdk-lib';
+import { Stack } from 'aws-cdk-lib';
 import {
-  // biome-ignore lint/suspicious/noShadowRestrictedNames: <explanation>
+  Architecture,
+  // biome-ignore lint/suspicious/noShadowRestrictedNames: shadows 'function'
   Function,
   type FunctionOptions,
   Runtime,
@@ -11,12 +11,19 @@ import {
 
 import type { Construct } from 'constructs';
 import { Bundling } from './bundling';
+import type { BundlingOptions } from './types';
 
 export interface PythonFunctionProps extends FunctionOptions {
   /**
-   * Path where index and function dependencies can be found
+   * UV project root directory (workspace root)
    */
-  readonly entry: string;
+  readonly rootDir: string;
+
+  /**
+   * Optional UV project workspace, used to specify a specific package to be used
+   * as a Lambda Function entry.
+   */
+  readonly workspacePackage?: string;
 
   /**
    * The runtime
@@ -26,7 +33,7 @@ export interface PythonFunctionProps extends FunctionOptions {
   readonly runtime?: Runtime;
 
   /**
-   * The path to the index file containing the handler. Relative to #entry
+   * The path to the index file with the project or (or workspace, if specified) containing the handler.
    *
    * @default index.py
    */
@@ -48,33 +55,34 @@ export interface PythonFunctionProps extends FunctionOptions {
 export class PythonFunction extends Function {
   constructor(scope: Construct, id: string, props: PythonFunctionProps) {
     const {
-      index = 'index.py',
+      workspacePackage,
       handler = 'handler',
       runtime = Runtime.PYTHON_3_12,
     } = props;
 
-    const entry = path.resolve(props.entry);
-    const resolvedIndex = path.resolve(entry, index);
-    if (!fs.existsSync(resolvedIndex)) {
-      throw new Error(`Cannot find index file at ${resolvedIndex}`);
-    }
+    const architecture = props.architecture ?? Architecture.ARM_64;
+    const rootDir = path.resolve(props.rootDir);
 
-    const moduleName = path.parse(index).name;
-    const resolvedHandler = `${moduleName}.${handler}`.replace(/\//g, '.');
+    let resolvedHandler = handler;
+    if (workspacePackage) {
+      resolvedHandler = `${workspacePackage.replace(/-/g, '_')}.${handler}`;
+    }
 
     if (runtime.family !== RuntimeFamily.PYTHON) {
       throw new Error('Only Python runtimes are supported');
     }
 
     const code = Bundling.bundle({
-      entry,
+      rootDir,
       runtime,
       skip: !Stack.of(scope).bundlingRequired,
-      architecture: props.architecture,
+      architecture,
+      workspacePackage,
       ...props.bundling,
     });
     super(scope, id, {
       ...props,
+      architecture,
       runtime,
       code,
       handler: resolvedHandler,

src/function.ts

@@ -1,2 +1,3 @@
 export * from './bundling';
 export * from './function';
+export * from './types';