Add patch_env to Python workers sdk

A context manager, to be used like:

```
with patch_env(key1=value1, key2=value2):
    # env is patched
# env is restored
```

Author: hoodmane

.gitignore

@@ -40,3 +40,6 @@ coverage
 perf.data
 
 .claude
+
+# python
+__pycache__

src/cloudflare/internal/env.d.ts

@@ -8,3 +8,6 @@ export function withEnvAndExports(
   newExports: unknown,
   fn: () => unknown
 ): unknown;
+export function patchEnv(newEnv: unknown): {
+  [Symbol.dispose](): void;
+};

src/pyodide/internal/envHelpers.ts

@@ -0,0 +1,13 @@
+// TODO: figure out how to make bazel resolve env.d.ts
+// eslint-disable-next-line
+// @ts-ignore
+import innerEnv from 'cloudflare-internal:env';
+
+// A generator that keeps the environment patched until it is resolved.
+// Used to define the patch_env context manager, see
+// src/pyodide/internal/workers-api/src/workers/_workers.py
+export function* patch_env_helper(patch: unknown): Generator<void> {
+  // eslint-disable-next-line
+  using _ = innerEnv.patchEnv(patch);
+  yield;
+}

src/pyodide/internal/snapshot.ts

@@ -21,6 +21,7 @@ import {
   PythonWorkersInternalError,
   PythonUserError,
   simpleRunPython,
+  unreachable,
 } from 'pyodide-internal:util';
 import { default as MetadataReader } from 'pyodide-internal:runtime-generated/metadata';
 import type { PyodideEntrypointHelper } from 'pyodide:python-entrypoint-helper';
@@ -556,6 +557,8 @@ ${describeValue(obj)}
   return new PythonUserError(error);
 }
 
+type CustomSerialized = { pyodide_entrypoint_helper: true } | { env: true };
+
 /**
  * Create memory snapshot by importing SNAPSHOT_IMPORTS to ensure these packages
  * are initialized in the linear memory snapshot and then saving a copy of the
@@ -567,10 +570,13 @@ function makeLinearMemorySnapshot(
   pyodide_entrypoint_helper: PyodideEntrypointHelper | null,
   snapshotType: ArtifactBundler.SnapshotType
 ): Uint8Array {
-  const customHiwireStateSerializer = (obj: any): Record<string, boolean> => {
+  const customHiwireStateSerializer = (obj: any): CustomSerialized => {
     if (obj === pyodide_entrypoint_helper) {
       return { pyodide_entrypoint_helper: true };
     }
+    if (obj === pyodide_entrypoint_helper?.cloudflareWorkersModule.env) {
+      return { env: true };
+    }
     throw createUnserializableObjectError(obj);
   };
 
@@ -843,11 +849,14 @@ export function finalizeBootstrap(
   Module: Module,
   pyodide_entrypoint_helper: PyodideEntrypointHelper
 ): void {
-  const customHiwireStateDeserializer = (obj: any): any => {
+  const customHiwireStateDeserializer = (obj: CustomSerialized): any => {
     if ('pyodide_entrypoint_helper' in obj) {
       return pyodide_entrypoint_helper;
     }
-    throw new PythonWorkersInternalError(`Can't deserialize ${obj}`);
+    if ('env' in obj) {
+      return pyodide_entrypoint_helper.cloudflareWorkersModule.env;
+    }
+    unreachable(obj, `Can't deserialize ${obj}`);
   };
 
   Module.API.config._makeSnapshot =

src/pyodide/internal/util.ts

@@ -79,6 +79,12 @@ export function invalidateCaches(Module: Module): void {
   );
 }
 
-export function unreachable(msg: never): never {
+export function unreachable(
+  obj: never,
+  msg: string | undefined = undefined
+): never {
+  if (msg === undefined) {
+    msg = obj;
+  }
   throw new PythonWorkersInternalError(`Unreachable: ${msg}`);
 }

src/pyodide/internal/workers-api/src/workers/__init__.py

@@ -20,6 +20,7 @@
     fetch,
     handler,
     import_from_javascript,
+    patch_env,
     python_from_rpc,
     python_to_rpc,
 )
@@ -47,6 +48,7 @@
     "fetch",
     "handler",
     "import_from_javascript",
+    "patch_env",
     "python_from_rpc",
     "python_to_rpc",
 ]

src/pyodide/internal/workers-api/src/workers/_workers.py

@@ -6,7 +6,14 @@
 import inspect
 import json
 from asyncio import create_task, gather
-from collections.abc import Awaitable, Generator, Iterable, MutableMapping
+from collections.abc import (
+    Awaitable,
+    Generator,
+    Iterable,
+    Iterator,
+    MutableMapping,
+    Sequence,
+)
 from contextlib import ExitStack, contextmanager
 from enum import StrEnum
 from http import HTTPMethod, HTTPStatus
@@ -96,6 +103,15 @@ def import_from_javascript(module_name: str) -> Any:
         raise
 
 
+@contextmanager
+def patch_env(
+    d: dict[str, Any] | Sequence[tuple[str, Any]] | None = None, **kwds: dict[str, Any]
+) -> Iterator[None]:
+    if d:
+        kwds = dict(d) | kwds
+    yield from _pyodide_entrypoint_helper.patch_env_helper(to_js(kwds))
+
+
 type JSBody = (
     "js.Blob | JsBuffer | js.FormData | js.ReadableStream | js.URLSearchParams"
 )

src/pyodide/python-entrypoint-helper.ts

@@ -27,6 +27,7 @@ import {
   reportError,
 } from 'pyodide-internal:util';
 import { LOADED_SNAPSHOT_TYPE } from 'pyodide-internal:snapshot';
+import { patch_env_helper } from 'pyodide-internal:envHelpers';
 
 type PyFuture<T> = Promise<T> & { copy(): PyFuture<T>; destroy(): void };
 
@@ -65,10 +66,11 @@ function patchWaitUntil(ctx: {
 
 export type PyodideEntrypointHelper = {
   doAnImport: (mod: string) => Promise<any>;
-  cloudflareWorkersModule: any;
+  cloudflareWorkersModule: { env: any };
   cloudflareSocketsModule: any;
   workerEntrypoint: any;
   patchWaitUntil: typeof patchWaitUntil;
+  patch_env_helper: (patch: unknown) => Generator<void>;
 };
 import { maybeCollectDedicatedSnapshot } from 'pyodide-internal:snapshot';
 
@@ -96,6 +98,7 @@ export function setDoAnImport(
     cloudflareSocketsModule,
     workerEntrypoint,
     patchWaitUntil,
+    patch_env_helper,
   };
 }
 

src/pyodide/tsconfig.json

@@ -10,7 +10,8 @@
       "pyodide-internal:generated/emscriptenSetup": [
         "./internal/pool/emscriptenSetup.ts"
       ],
-      "internal:*": ["./types/*"]
+      "internal:*": ["./types/*"],
+      "cloudflare-internal:*": ["../cloudflare/internal/*"]
     },
     "typeRoots": ["./types"],
     "noEmit": false,

src/workerd/api/modules.c++

@@ -37,6 +37,11 @@ jsg::JsRef<jsg::JsValue> EnvModule::withEnv(
   });
 }
 
+jsg::Ref<PatchedEnv> EnvModule::patchEnv(jsg::Lock& js, jsg::Value newEnv) {
+  auto& key = jsg::IsolateBase::from(js.v8Isolate).getEnvAsyncContextKey();
+  return jsg::alloc<PatchedEnv>(js, key, kj::mv(newEnv));
+}
+
 kj::Maybe<jsg::JsObject> EnvModule::getCurrentExports(jsg::Lock& js) {
   auto& key = jsg::IsolateBase::from(js.v8Isolate).getExportsAsyncContextKey();
   // Check async context first - withExports() overrides take precedence over the disable flags.