pythonExtension and python runtime improvements (#1734)

* pythonExtension and python runtime improvements

* Adding streaming support

* Use writeFileSync

* Restructure extension docs and add python extension docs

* Fix broken link

* Update docs/config/extensions/overview.mdx

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>

* Update docs/config/extensions/aptGet.mdx

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>

* Update docs/config/extensions/custom.mdx

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>

* Add environment variable support

---------

Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>

packages/python/README.md

@@ -14,7 +14,7 @@ This extension introduces the <code>pythonExtension</code> build extension, whic
   - <code>run</code>: Executes Python commands with proper environment setup.
   - <code>runInline</code>: Executes inline Python code directly from Node.
   - <code>runScript</code>: Executes standalone <code>.py</code> script files.
-- **Custom Python Path:** In development, you can configure <code>pythonBinaryPath</code> to point to a custom Python installation.
+- **Custom Python Path:** In development, you can configure <code>devPythonBinaryPath</code> to point to a custom Python installation.
 
 ## Usage
 
@@ -22,7 +22,7 @@ This extension introduces the <code>pythonExtension</code> build extension, whic
 
 ```typescript
 import { defineConfig } from "@trigger.dev/sdk/v3";
-import pythonExtension from "@trigger.dev/python/extension";
+import { pythonExtension } from "@trigger.dev/python/extension";
 
 export default defineConfig({
   project: "<project ref>",
@@ -30,8 +30,8 @@ export default defineConfig({
     extensions: [
       pythonExtension({
         requirementsFile: "./requirements.txt", // Optional: Path to your requirements file
-        pythonBinaryPath: path.join(rootDir, `.venv/bin/python`), // Optional: Custom Python binary path
-        scripts: ["my_script.py"], // List of Python scripts to include
+        devPythonBinaryPath: ".venv/bin/python", // Optional: Custom Python binary path
+        scripts: ["src/python/**/*.py"], // Glob pattern for Python scripts
       }),
     ],
   },
@@ -40,13 +40,34 @@ export default defineConfig({
 
 2. (Optional) Create a <code>requirements.txt</code> file in your project root with the necessary Python dependencies.
 
+```plaintext title="requirements.txt"
+pandas==1.3.3
+numpy==1.21.2
+```
+
+```typescript title="trigger.config.ts"
+import { defineConfig } from "@trigger.dev/sdk/v3";
+import { pythonExtension } from "@trigger.dev/python/extension";
+
+export default defineConfig({
+  project: "<project ref>",
+  build: {
+    extensions: [
+      pythonExtension({
+        requirementsFile: "./requirements.txt",
+      }),
+    ],
+  },
+});
+```
+
 3. Execute Python scripts within your tasks using one of the provided functions:
 
 ### Running a Python Script
 
 ```typescript
 import { task } from "@trigger.dev/sdk/v3";
-import python from "@trigger.dev/python";
+import { python } from "@trigger.dev/python";
 
 export const myScript = task({
   id: "my-python-script",
@@ -55,13 +76,29 @@ export const myScript = task({
     return result.stdout;
   },
 });
+
+export const myStreamingScript = task({
+  id: "my-streaming-python-script",
+  run: async () => {
+    // You can also stream the output of the script
+    const result = python.stream.runScript("my_script.py", ["hello", "world"]);
+
+    // result is an async iterable/readable stream
+    for await (const chunk of streamingResult) {
+      logger.debug("convert-url-to-markdown", {
+        url: payload.url,
+        chunk,
+      });
+    }
+  },
+});
 ```
 
 ### Running Inline Python Code
 
 ```typescript
 import { task } from "@trigger.dev/sdk/v3";
-import python from "@trigger.dev/python";
+import { python } from "@trigger.dev/python";
 
 export const myTask = task({
   id: "to_datetime-task",
@@ -69,7 +106,7 @@ export const myTask = task({
     const result = await python.runInline(`
 import pandas as pd
 
-pandas.to_datetime("${+new Date() / 1000}")
+pd.to_datetime("${+new Date() / 1000}")
 `);
     return result.stdout;
   },
@@ -80,7 +117,7 @@ pandas.to_datetime("${+new Date() / 1000}")
 
 ```typescript
 import { task } from "@trigger.dev/sdk/v3";
-import python from "@trigger.dev/python";
+import { python } from "@trigger.dev/python";
 
 export const pythonVersionTask = task({
   id: "python-version-task",
@@ -94,7 +131,6 @@ export const pythonVersionTask = task({
 ## Limitations
 
 - This is a **partial implementation** and does not provide full Python support as an execution runtime for tasks.
-- Only basic Python script execution is supported; scripts are not automatically copied to staging/production containers.
 - Manual intervention may be required for installing and configuring binary dependencies in development environments.
 
 ## Additional Information

packages/python/package.json

@@ -45,9 +45,7 @@
     "check-exports": "attw --pack ."
   },
   "dependencies": {
-    "@trigger.dev/build": "workspace:3.3.16",
     "@trigger.dev/core": "workspace:3.3.16",
-    "@trigger.dev/sdk": "workspace:3.3.16",
     "tinyexec": "^0.3.2"
   },
   "devDependencies": {
@@ -57,7 +55,13 @@
     "typescript": "^5.5.4",
     "tsx": "4.17.0",
     "esbuild": "^0.23.0",
-    "@arethetypeswrong/cli": "^0.15.4"
+    "@arethetypeswrong/cli": "^0.15.4",
+    "@trigger.dev/build": "workspace:3.3.16",
+    "@trigger.dev/sdk": "workspace:3.3.16"
+  },
+  "peerDependencies": {
+    "@trigger.dev/sdk": "workspace:^3.3.16",
+    "@trigger.dev/build": "workspace:^3.3.16"
   },
   "engines": {
     "node": ">=18.20.0"

packages/python/src/extension.ts

@@ -1,6 +1,6 @@
 import fs from "node:fs";
 import assert from "node:assert";
-import { additionalFiles } from "@trigger.dev/build/extensions/core";
+import { addAdditionalFilesToBuild } from "@trigger.dev/build/internal";
 import { BuildManifest } from "@trigger.dev/core/v3";
 import { BuildContext, BuildExtension } from "@trigger.dev/core/v3/build";
 
@@ -16,7 +16,7 @@ export type PythonOptions = {
    *
    * Example: `/usr/bin/python3` or `C:\\Python39\\python.exe`
    */
-  pythonBinaryPath?: string;
+  devPythonBinaryPath?: string;
   /**
    * An array of glob patterns that specify which Python scripts are allowed to be executed.
    *
@@ -57,13 +57,18 @@ class PythonExtension implements BuildExtension {
   }
 
   async onBuildComplete(context: BuildContext, manifest: BuildManifest) {
-    await additionalFiles({
-      files: this.options.scripts ?? [],
-    }).onBuildComplete!(context, manifest);
+    await addAdditionalFilesToBuild(
+      "pythonExtension",
+      {
+        files: this.options.scripts ?? [],
+      },
+      context,
+      manifest
+    );
 
     if (context.target === "dev") {
-      if (this.options.pythonBinaryPath) {
-        process.env.PYTHON_BIN_PATH = this.options.pythonBinaryPath;
+      if (this.options.devPythonBinaryPath) {
+        process.env.PYTHON_BIN_PATH = this.options.devPythonBinaryPath;
       }
 
       return;
@@ -93,27 +98,59 @@ class PythonExtension implements BuildExtension {
       },
     });
 
-    context.addLayer({
-      id: "python-dependencies",
-      build: {
-        env: {
-          REQUIREMENTS_CONTENT: this.options.requirements?.join("\n") || "",
+    if (this.options.requirementsFile) {
+      if (this.options.requirements) {
+        context.logger.warn(
+          `[pythonExtension] Both options.requirements and options.requirementsFile are specified. requirements will be ignored.`
+        );
+      }
+
+      // Copy requirements file to the container
+      await addAdditionalFilesToBuild(
+        "pythonExtension",
+        {
+          files: [this.options.requirementsFile],
         },
-      },
-      image: {
-        instructions: splitAndCleanComments(`
+        context,
+        manifest
+      );
+
+      // Add a layer to the build that installs the requirements
+      context.addLayer({
+        id: "python-dependencies",
+        image: {
+          instructions: splitAndCleanComments(`
+            # Copy the requirements file
+            COPY ${this.options.requirementsFile} .
+            # Install dependencies
+            RUN pip install --no-cache-dir -r ${this.options.requirementsFile}
+          `),
+        },
+        deploy: {
+          override: true,
+        },
+      });
+    } else if (this.options.requirements) {
+      context.addLayer({
+        id: "python-dependencies",
+        build: {
+          env: {
+            REQUIREMENTS_CONTENT: this.options.requirements?.join("\n") || "",
+          },
+        },
+        image: {
+          instructions: splitAndCleanComments(`
           ARG REQUIREMENTS_CONTENT
           RUN echo "$REQUIREMENTS_CONTENT" > requirements.txt
 
           # Install dependencies
           RUN pip install --no-cache-dir -r requirements.txt
         `),
-      },
-      deploy: {
-        override: true,
-      },
-    });
+        },
+        deploy: {
+          override: true,
+        },
+      });
+    }
   }
 }
-
-export default pythonExtension;

packages/python/src/index.ts

@@ -1,63 +1,277 @@
-import fs from "node:fs";
-import assert from "node:assert";
+import {
+  AsyncIterableStream,
+  createAsyncIterableStreamFromAsyncIterable,
+  SemanticInternalAttributes,
+} from "@trigger.dev/core/v3";
 import { logger } from "@trigger.dev/sdk/v3";
-import { x, Options as XOptions, Result } from "tinyexec";
+import assert from "node:assert";
+import fs from "node:fs";
+import { Result, x, Options as XOptions } from "tinyexec";
+import { createTempFileSync, withTempFile } from "./utils/tempFiles.js";
 
-export const run = async (
-  scriptArgs: string[] = [],
-  options: Partial<XOptions> = {}
-): Promise<Result> => {
-  const pythonBin = process.env.PYTHON_BIN_PATH || "python";
-
-  return await logger.trace("Python call", async (span) => {
-    span.addEvent("Properties", {
-      command: `${pythonBin} ${scriptArgs.join(" ")}`,
-    });
-
-    const result = await x(pythonBin, scriptArgs, {
-      ...options,
-      throwOnError: false, // Ensure errors are handled manually
-    });
-
-    span.addEvent("Output", { ...result });
-
-    if (result.exitCode !== 0) {
-      logger.error(result.stderr, { ...result });
-      throw new Error(`Python command exited with non-zero code ${result.exitCode}`);
-    }
-
-    return result;
-  });
+export type PythonExecOptions = Partial<XOptions> & {
+  env?: { [key: string]: string | undefined };
 };
 
-export const runScript = (
-  scriptPath: string,
-  scriptArgs: string[] = [],
-  options: Partial<XOptions> = {}
-) => {
-  assert(scriptPath, "Script path is required");
-  assert(fs.existsSync(scriptPath), `Script does not exist: ${scriptPath}`);
+export const python = {
+  async run(scriptArgs: string[] = [], options: PythonExecOptions = {}): Promise<Result> {
+    const pythonBin = process.env.PYTHON_BIN_PATH || "python";
 
-  return run([scriptPath, ...scriptArgs], options);
-};
+    return await logger.trace(
+      "python.run()",
+      async (span) => {
+        const result = await x(pythonBin, scriptArgs, {
+          ...options,
+          nodeOptions: {
+            ...(options.nodeOptions || {}),
+            env: {
+              ...process.env,
+              ...options.env,
+            },
+          },
+          throwOnError: false, // Ensure errors are handled manually
+        });
 
-export const runInline = async (scriptContent: string, options: Partial<XOptions> = {}) => {
-  assert(scriptContent, "Script content is required");
+        if (result.exitCode) {
+          span.setAttribute("exitCode", result.exitCode);
+        }
 
-  const tmpFile = `/tmp/script_${Date.now()}.py`;
-  await fs.promises.writeFile(tmpFile, scriptContent, { mode: 0o600 });
+        if (result.exitCode !== 0) {
+          throw new Error(
+            `${scriptArgs.join(" ")} exited with a non-zero code ${result.exitCode}:\n${
+              result.stderr
+            }`
+          );
+        }
 
-  try {
-    return await runScript(tmpFile, [], options);
-  } finally {
-    try {
-      await fs.promises.unlink(tmpFile);
-    } catch (error) {
-      logger.warn(`Failed to clean up temporary file ${tmpFile}:`, {
-        error: (error as Error).stack || (error as Error).message,
+        return result;
+      },
+      {
+        attributes: {
+          pythonBin,
+          args: scriptArgs.join(" "),
+          [SemanticInternalAttributes.STYLE_ICON]: "brand-python",
+        },
+      }
+    );
+  },
+
+  async runScript(
+    scriptPath: string,
+    scriptArgs: string[] = [],
+    options: PythonExecOptions = {}
+  ): Promise<Result> {
+    assert(scriptPath, "Script path is required");
+    assert(fs.existsSync(scriptPath), `Script does not exist: ${scriptPath}`);
+
+    return await logger.trace(
+      "python.runScript()",
+      async (span) => {
+        span.setAttribute("scriptPath", scriptPath);
+
+        const result = await x(
+          process.env.PYTHON_BIN_PATH || "python",
+          [scriptPath, ...scriptArgs],
+          {
+            ...options,
+            nodeOptions: {
+              ...(options.nodeOptions || {}),
+              env: {
+                ...process.env,
+                ...options.env,
+              },
+            },
+            throwOnError: false,
+          }
+        );
+
+        if (result.exitCode) {
+          span.setAttribute("exitCode", result.exitCode);
+        }
+
+        if (result.exitCode !== 0) {
+          throw new Error(
+            `${scriptPath} ${scriptArgs.join(" ")} exited with a non-zero code ${
+              result.exitCode
+            }:\n${result.stderr}`
+          );
+        }
+
+        return result;
+      },
+      {
+        attributes: {
+          pythonBin: process.env.PYTHON_BIN_PATH || "python",
+          scriptPath,
+          args: scriptArgs.join(" "),
+          [SemanticInternalAttributes.STYLE_ICON]: "brand-python",
+        },
+      }
+    );
+  },
+
+  async runInline(scriptContent: string, options: PythonExecOptions = {}): Promise<Result> {
+    assert(scriptContent, "Script content is required");
+
+    return await logger.trace(
+      "python.runInline()",
+      async (span) => {
+        span.setAttribute("contentLength", scriptContent.length);
+
+        // Using the withTempFile utility to handle the temporary file
+        return await withTempFile(
+          `script_${Date.now()}.py`,
+          async (tempFilePath) => {
+            span.setAttribute("tempFilePath", tempFilePath);
+
+            const pythonBin = process.env.PYTHON_BIN_PATH || "python";
+            const result = await x(pythonBin, [tempFilePath], {
+              ...options,
+              nodeOptions: {
+                ...(options.nodeOptions || {}),
+                env: {
+                  ...process.env,
+                  ...options.env,
+                },
+              },
+              throwOnError: false,
+            });
+
+            if (result.exitCode) {
+              span.setAttribute("exitCode", result.exitCode);
+            }
+
+            if (result.exitCode !== 0) {
+              throw new Error(
+                `Inline script exited with a non-zero code ${result.exitCode}:\n${result.stderr}`
+              );
+            }
+
+            return result;
+          },
+          scriptContent
+        );
+      },
+      {
+        attributes: {
+          pythonBin: process.env.PYTHON_BIN_PATH || "python",
+          contentPreview:
+            scriptContent.substring(0, 100) + (scriptContent.length > 100 ? "..." : ""),
+          [SemanticInternalAttributes.STYLE_ICON]: "brand-python",
+        },
+      }
+    );
+  },
+  // Stream namespace for streaming functions
+  stream: {
+    run(scriptArgs: string[] = [], options: PythonExecOptions = {}): AsyncIterableStream<string> {
+      const pythonBin = process.env.PYTHON_BIN_PATH || "python";
+
+      const pythonProcess = x(pythonBin, scriptArgs, {
+        ...options,
+        nodeOptions: {
+          ...(options.nodeOptions || {}),
+          env: {
+            ...process.env,
+            ...options.env,
+          },
+        },
+        throwOnError: false,
       });
-    }
-  }
-};
 
-export default { run, runScript, runInline };
+      const span = logger.startSpan("python.stream.run()", {
+        attributes: {
+          pythonBin,
+          args: scriptArgs.join(" "),
+          [SemanticInternalAttributes.STYLE_ICON]: "brand-python",
+        },
+      });
+
+      return createAsyncIterableStreamFromAsyncIterable(pythonProcess, {
+        transform: (chunk, controller) => {
+          controller.enqueue(chunk);
+        },
+        flush: () => {
+          span.end();
+        },
+      });
+    },
+    runScript(
+      scriptPath: string,
+      scriptArgs: string[] = [],
+      options: PythonExecOptions = {}
+    ): AsyncIterableStream<string> {
+      assert(scriptPath, "Script path is required");
+      assert(fs.existsSync(scriptPath), `Script does not exist: ${scriptPath}`);
+
+      const pythonBin = process.env.PYTHON_BIN_PATH || "python";
+
+      const pythonProcess = x(pythonBin, [scriptPath, ...scriptArgs], {
+        ...options,
+        nodeOptions: {
+          ...(options.nodeOptions || {}),
+          env: {
+            ...process.env,
+            ...options.env,
+          },
+        },
+        throwOnError: false,
+      });
+
+      const span = logger.startSpan("python.stream.runScript()", {
+        attributes: {
+          pythonBin,
+          scriptPath,
+          args: scriptArgs.join(" "),
+          [SemanticInternalAttributes.STYLE_ICON]: "brand-python",
+        },
+      });
+
+      return createAsyncIterableStreamFromAsyncIterable(pythonProcess, {
+        transform: (chunk, controller) => {
+          controller.enqueue(chunk);
+        },
+        flush: () => {
+          span.end();
+        },
+      });
+    },
+    runInline(scriptContent: string, options: PythonExecOptions = {}): AsyncIterableStream<string> {
+      assert(scriptContent, "Script content is required");
+
+      const pythonBin = process.env.PYTHON_BIN_PATH || "python";
+
+      const pythonScriptPath = createTempFileSync(`script_${Date.now()}.py`, scriptContent);
+
+      const pythonProcess = x(pythonBin, [pythonScriptPath], {
+        ...options,
+        nodeOptions: {
+          ...(options.nodeOptions || {}),
+          env: {
+            ...process.env,
+            ...options.env,
+          },
+        },
+        throwOnError: false,
+      });
+
+      const span = logger.startSpan("python.stream.runInline()", {
+        attributes: {
+          pythonBin,
+          contentPreview:
+            scriptContent.substring(0, 100) + (scriptContent.length > 100 ? "..." : ""),
+          [SemanticInternalAttributes.STYLE_ICON]: "brand-python",
+        },
+      });
+
+      return createAsyncIterableStreamFromAsyncIterable(pythonProcess, {
+        transform: (chunk, controller) => {
+          controller.enqueue(chunk);
+        },
+        flush: () => {
+          span.end();
+        },
+      });
+    },
+  },
+};

packages/python/src/utils/tempFiles.ts

@@ -0,0 +1,39 @@
+import { mkdtempSync, writeFileSync } from "node:fs";
+import { mkdtemp, writeFile, rm } from "node:fs/promises";
+import { tmpdir } from "node:os";
+import { join } from "node:path";
+
+/**
+ * Creates a temporary file with a custom filename, passes it to the callback function, and ensures cleanup
+ * @param filename The filename to use for the temporary file
+ * @param callback Function that receives the path to the temporary file
+ * @param content Optional content to write to the file
+ * @returns Whatever the callback returns
+ */
+export async function withTempFile<T>(
+  filename: string,
+  callback: (filePath: string) => Promise<T>,
+  content: string | Buffer = ""
+): Promise<T> {
+  // Create temporary directory with random suffix
+  const tempDir = await mkdtemp(join(tmpdir(), "app-"));
+  const tempFile = join(tempDir, filename);
+
+  try {
+    // Write to the temporary file with appropriate permissions
+    await writeFile(tempFile, content, { mode: 0o600 });
+    // Use the file
+    return await callback(tempFile);
+  } finally {
+    // Clean up
+    await rm(tempDir, { recursive: true, force: true });
+  }
+}
+
+export function createTempFileSync(filename: string, content: string | Buffer = ""): string {
+  const tempDir = mkdtempSync(join(tmpdir(), "app-"));
+  const tempFile = join(tempDir, filename);
+
+  writeFileSync(tempFile, content, { mode: 0o600 });
+  return tempFile;
+}