Add environment variable support

Author: ericallam

docs/config/extensions/pythonExtension.mdx

@@ -151,3 +151,32 @@ export const myStreamingScript = task({
   },
 });
 ```
+
+## Environment variables
+
+We automatically inject the environment variables in the `process.env` object when running Python scripts. You can access these environment variables in your Python scripts using the `os.environ` dictionary. For example:
+
+```python
+import os
+
+print(os.environ["MY_ENV_VAR"])
+```
+
+You can also pass additional environment variables to the Python script using the `env` option in the `python.runScript` function. For example:
+
+```ts
+import { task } from "@trigger.dev/sdk/v3";
+import { python } from "@trigger.dev/python";
+
+export const myScript = task({
+  id: "my-python-script",
+  run: async () => {
+    const result = await python.runScript("./python/my_script.py", ["hello", "world"], {
+      env: {
+        MY_ENV_VAR: "my value",
+      },
+    });
+    return result.stdout;
+  },
+});
+```

packages/python/src/index.ts

@@ -9,15 +9,26 @@ import fs from "node:fs";
 import { Result, x, Options as XOptions } from "tinyexec";
 import { createTempFileSync, withTempFile } from "./utils/tempFiles.js";
 
+export type PythonExecOptions = Partial<XOptions> & {
+  env?: { [key: string]: string | undefined };
+};
+
 export const python = {
-  async run(scriptArgs: string[] = [], options: Partial<XOptions> = {}): Promise<Result> {
+  async run(scriptArgs: string[] = [], options: PythonExecOptions = {}): Promise<Result> {
     const pythonBin = process.env.PYTHON_BIN_PATH || "python";
 
     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
         });
 
@@ -48,7 +59,7 @@ export const python = {
   async runScript(
     scriptPath: string,
     scriptArgs: string[] = [],
-    options: Partial<XOptions> = {}
+    options: PythonExecOptions = {}
   ): Promise<Result> {
     assert(scriptPath, "Script path is required");
     assert(fs.existsSync(scriptPath), `Script does not exist: ${scriptPath}`);
@@ -63,6 +74,13 @@ export const python = {
           [scriptPath, ...scriptArgs],
           {
             ...options,
+            nodeOptions: {
+              ...(options.nodeOptions || {}),
+              env: {
+                ...process.env,
+                ...options.env,
+              },
+            },
             throwOnError: false,
           }
         );
@@ -92,7 +110,7 @@ export const python = {
     );
   },
 
-  async runInline(scriptContent: string, options: Partial<XOptions> = {}): Promise<Result> {
+  async runInline(scriptContent: string, options: PythonExecOptions = {}): Promise<Result> {
     assert(scriptContent, "Script content is required");
 
     return await logger.trace(
@@ -109,6 +127,13 @@ export const python = {
             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,
             });
 
@@ -139,11 +164,18 @@ export const python = {
   },
   // Stream namespace for streaming functions
   stream: {
-    run(scriptArgs: string[] = [], options: Partial<XOptions> = {}): AsyncIterableStream<string> {
+    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,
       });
 
@@ -167,7 +199,7 @@ export const python = {
     runScript(
       scriptPath: string,
       scriptArgs: string[] = [],
-      options: Partial<XOptions> = {}
+      options: PythonExecOptions = {}
     ): AsyncIterableStream<string> {
       assert(scriptPath, "Script path is required");
       assert(fs.existsSync(scriptPath), `Script does not exist: ${scriptPath}`);
@@ -176,6 +208,13 @@ export const python = {
 
       const pythonProcess = x(pythonBin, [scriptPath, ...scriptArgs], {
         ...options,
+        nodeOptions: {
+          ...(options.nodeOptions || {}),
+          env: {
+            ...process.env,
+            ...options.env,
+          },
+        },
         throwOnError: false,
       });
 
@@ -197,7 +236,7 @@ export const python = {
         },
       });
     },
-    runInline(scriptContent: string, options: Partial<XOptions> = {}): AsyncIterableStream<string> {
+    runInline(scriptContent: string, options: PythonExecOptions = {}): AsyncIterableStream<string> {
       assert(scriptContent, "Script content is required");
 
       const pythonBin = process.env.PYTHON_BIN_PATH || "python";
@@ -206,6 +245,13 @@ export const python = {
 
       const pythonProcess = x(pythonBin, [pythonScriptPath], {
         ...options,
+        nodeOptions: {
+          ...(options.nodeOptions || {}),
+          env: {
+            ...process.env,
+            ...options.env,
+          },
+        },
         throwOnError: false,
       });
 

references/python-catalog/src/trigger/pythonTasks.ts

@@ -29,13 +29,24 @@ export const convertUrlToMarkdown = schemaTask({
 export const pythonRunInlineTask = task({
   id: "python-run-inline",
   run: async () => {
-    const result = await python.runInline(`
+    const result = await python.runInline(
+      `
+import os
 import html2text as h2t
 
 h = h2t.HTML2Text()
 
 print(h.handle("<p>Hello, <a href='https://www.google.com/earth/'>world</a>!"))
-`);
+print(f"API Key: {os.environ['OPENAI_API_KEY']}")
+`,
+      {
+        env: {
+          OPENAI_API_KEY: "sk-1234567890",
+        },
+      }
+    );
+
+    console.log(result.stdout);
 
     const streamingResult = python.stream.runInline(`
 import html2text as h2t