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/pyodide/internal/envHelpers.ts

@@ -0,0 +1,9 @@
+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> {
+  using _ = innerEnv.pythonPatchEnv(patch);
+  yield;
+}

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:*": ["./types/cloudflare-internal/*"]
     },
     "typeRoots": ["./types"],
     "noEmit": false,

src/pyodide/types/cloudflare-internal/env.d.ts

@@ -0,0 +1,6 @@
+const innerEnv: {
+  pythonPatchEnv(patch: unknown): {
+    [Symbol.dispose](): void;
+  };
+};
+export default innerEnv;

src/workerd/api/modules.c++

@@ -37,6 +37,11 @@ jsg::JsRef<jsg::JsValue> EnvModule::withEnv(
   });
 }
 
+jsg::Ref<PythonPatchedEnv> EnvModule::pythonPatchEnv(jsg::Lock& js, jsg::Value newEnv) {
+  auto& key = jsg::IsolateBase::from(js.v8Isolate).getEnvAsyncContextKey();
+  return jsg::alloc<PythonPatchedEnv>(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.

src/workerd/api/modules.h

@@ -19,6 +19,27 @@
 
 namespace workerd::api {
 
+// An object with a [Symbol.dispose]() method to remove patch to environment. Not exposed
+// publically, just used to implement Python's `patch_env()` context manager.
+// See src/pyodide/internal/envHelpers.ts
+class PythonPatchedEnv: public jsg::Object {
+ public:
+  PythonPatchedEnv(jsg::Lock& js, jsg::AsyncContextFrame::StorageKey& key, jsg::Value store) {
+    scope.emplace(js, key, kj::mv(store));
+  }
+
+  void dispose() {
+    scope = kj::none;
+  }
+
+  JSG_RESOURCE_TYPE(PythonPatchedEnv) {
+    JSG_DISPOSE(dispose);
+  }
+
+ private:
+  kj::Maybe<jsg::AsyncContextFrame::StorageScope> scope;
+};
+
 class EnvModule final: public jsg::Object {
  public:
   EnvModule() = default;
@@ -39,12 +60,18 @@ class EnvModule final: public jsg::Object {
       jsg::Value newExports,
       jsg::Function<jsg::JsRef<jsg::JsValue>()> fn);
 
+  // Patch environment and return an object with a [Symbol.dispose]() method to restore it.
+  // Not exposed publically, just used to implement Python's `patch_env()` context manager.
+  // See src/pyodide/internal/envHelpers.ts
+  jsg::Ref<PythonPatchedEnv> pythonPatchEnv(jsg::Lock& js, jsg::Value newEnv);
+
   JSG_RESOURCE_TYPE(EnvModule) {
     JSG_METHOD(getCurrentEnv);
     JSG_METHOD(getCurrentExports);
     JSG_METHOD(withEnv);
     JSG_METHOD(withExports);
     JSG_METHOD(withEnvAndExports);
+    JSG_METHOD(pythonPatchEnv);
   }
 };