Allow importing python compat flag from python workers

Author: ryanking13

src/pyodide/internal/metadata.ts

@@ -47,8 +47,12 @@ export const TRANSITIVE_REQUIREMENTS =
 export const MAIN_MODULE_NAME = MetadataReader.getMainModule();
 
 export type CompatibilityFlags = MetadataReader.CompatibilityFlags;
-export const COMPATIBILITY_FLAGS: MetadataReader.CompatibilityFlags =
-  MetadataReader.getCompatibilityFlags();
+export const COMPATIBILITY_FLAGS: MetadataReader.CompatibilityFlags = {
+  // Compat flags returned from getCompatibilityFlags is immutable,
+  // but in Pyodide 0.26, we modify the JS object that is exposed to the Python through
+  // registerJsModule so we create a new object here by copying the values.
+  ...MetadataReader.getCompatibilityFlags(),
+};
 export const WORKFLOWS_ENABLED: boolean =
   !!COMPATIBILITY_FLAGS.python_workflows;
 const NO_GLOBAL_HANDLERS: boolean =

src/pyodide/internal/python.ts

@@ -8,6 +8,7 @@ import {
   maybeRestoreSnapshot,
   finalizeBootstrap,
   isRestoringSnapshot,
+  type CustomSerializedObjects,
 } from 'pyodide-internal:snapshot';
 import {
   entropyMountFiles,
@@ -20,7 +21,6 @@ import {
   LEGACY_VENDOR_PATH,
   setCpuLimitNearlyExceededCallback,
 } from 'pyodide-internal:metadata';
-import type { PyodideEntrypointHelper } from 'pyodide:python-entrypoint-helper';
 
 /**
  * SetupEmscripten is an internal module defined in setup-emscripten.h the module instantiates
@@ -47,7 +47,7 @@ import { TRANSITIVE_REQUIREMENTS } from 'pyodide-internal:metadata';
  */
 function prepareWasmLinearMemory(
   Module: Module,
-  pyodide_entrypoint_helper: PyodideEntrypointHelper
+  customSerializedObjects: CustomSerializedObjects
 ): void {
   maybeRestoreSnapshot(Module);
   // entropyAfterRuntimeInit adjusts JS state ==> always needs to be called.
@@ -61,7 +61,7 @@ function prepareWasmLinearMemory(
     adjustSysPath(Module);
   }
   if (Module.API.version !== '0.26.0a2') {
-    finalizeBootstrap(Module, pyodide_entrypoint_helper);
+    finalizeBootstrap(Module, customSerializedObjects);
   }
 }
 
@@ -210,7 +210,7 @@ export function loadPyodide(
   isWorkerd: boolean,
   lockfile: PackageLock,
   indexURL: string,
-  pyodide_entrypoint_helper: PyodideEntrypointHelper
+  customSerializedObjects: CustomSerializedObjects
 ): Pyodide {
   try {
     const Module = enterJaegerSpan('instantiate_emscripten', () =>
@@ -238,18 +238,18 @@ export function loadPyodide(
     });
 
     enterJaegerSpan('prepare_wasm_linear_memory', () => {
-      prepareWasmLinearMemory(Module, pyodide_entrypoint_helper);
+      prepareWasmLinearMemory(Module, customSerializedObjects);
     });
 
-    maybeCollectSnapshot(Module, pyodide_entrypoint_helper);
+    maybeCollectSnapshot(Module, customSerializedObjects);
     // Mount worker files after doing snapshot upload so we ensure that data from the files is never
     // present in snapshot memory.
     mountWorkerFiles(Module);
 
     if (Module.API.version === '0.26.0a2') {
       // Finish setting up Pyodide's ffi so we can use the nice Python interface
       // In newer versions we already did this in prepareWasmLinearMemory.
-      finalizeBootstrap(Module, pyodide_entrypoint_helper);
+      finalizeBootstrap(Module, customSerializedObjects);
     }
     const pyodide = Module.API.public_api;
 

src/pyodide/internal/snapshot.ts

@@ -582,15 +582,25 @@ async function importJsModulesFromSnapshot(
 
 type CustomSerialized =
   | { pyodide_entrypoint_helper: true }
+  | { cloudflare_compat_flags: true }
   | SerializedJsModule;
+/**
+ * Global objects that need a custom serializer
+ */
+export type CustomSerializedObjects = {
+  pyodide_entrypoint_helper: PyodideEntrypointHelper;
+  cloudflare_compat_flags: CompatibilityFlags;
+};
 
 function getHiwireSerializer(
-  pyodide_entrypoint_helper: PyodideEntrypointHelper,
+  globalObj: CustomSerializedObjects,
   modules: Set<string>
 ): (obj: any) => CustomSerialized {
   return function serializer(obj: any): CustomSerialized {
-    if (obj === pyodide_entrypoint_helper) {
+    if (obj === globalObj.pyodide_entrypoint_helper) {
       return { pyodide_entrypoint_helper: true };
+    } else if (obj === globalObj.cloudflare_compat_flags) {
+      return { cloudflare_compat_flags: true };
     }
     const serializedModule = maybeSerializeJsModule(obj, modules);
     if (serializedModule) {
@@ -601,11 +611,13 @@ function getHiwireSerializer(
 }
 
 function getHiwireDeserializer(
-  pyodide_entrypoint_helper: PyodideEntrypointHelper
+  globalObj: CustomSerializedObjects
 ): (obj: CustomSerialized) => any {
   return function deserializer(obj) {
     if ('pyodide_entrypoint_helper' in obj) {
-      return pyodide_entrypoint_helper;
+      return globalObj.pyodide_entrypoint_helper;
+    } else if ('cloudflare_compat_flags' in obj) {
+      return globalObj.cloudflare_compat_flags;
     }
     if ('jsModule' in obj) {
       return deserializeJsModule(obj, JS_MODULES);
@@ -622,15 +634,15 @@ function getHiwireDeserializer(
 function makeLinearMemorySnapshot(
   Module: Module,
   importedModulesList: string[],
-  pyodide_entrypoint_helper: PyodideEntrypointHelper,
+  customSerializedObjects: CustomSerializedObjects,
   snapshotType: ArtifactBundler.SnapshotType
 ): Uint8Array {
   const dsoHandles = recordDsoHandles(Module);
   let hiwire: SnapshotConfig | undefined;
   const jsModuleNames: Set<string> = new Set();
   if (Module.API.version !== '0.26.0a2') {
     hiwire = Module.API.serializeHiwireState(
-      getHiwireSerializer(pyodide_entrypoint_helper, jsModuleNames)
+      getHiwireSerializer(customSerializedObjects, jsModuleNames)
     );
   }
   const settings: SnapshotSettings = {
@@ -807,7 +819,7 @@ export function maybeRestoreSnapshot(Module: Module): void {
 function collectSnapshot(
   Module: Module,
   importedModulesList: string[],
-  pyodide_entrypoint_helper: PyodideEntrypointHelper,
+  customSerializedObjects: CustomSerializedObjects,
   snapshotType: ArtifactBundler.SnapshotType
 ): void {
   if (!IS_EW_VALIDATING && !SHOULD_SNAPSHOT_TO_DISK) {
@@ -818,7 +830,7 @@ function collectSnapshot(
   const snapshot = makeLinearMemorySnapshot(
     Module,
     importedModulesList,
-    pyodide_entrypoint_helper,
+    customSerializedObjects,
     snapshotType
   );
   entropyAfterSnapshot(Module);
@@ -841,7 +853,7 @@ function collectSnapshot(
  */
 export function maybeCollectDedicatedSnapshot(
   Module: Module,
-  pyodide_entrypoint_helper: PyodideEntrypointHelper | null
+  customSerializedObjects: CustomSerializedObjects | null
 ): void {
   if (!IS_CREATING_SNAPSHOT) {
     return;
@@ -859,12 +871,12 @@ export function maybeCollectDedicatedSnapshot(
     );
   }
 
-  if (!pyodide_entrypoint_helper) {
+  if (!customSerializedObjects) {
     throw new PythonWorkersInternalError(
-      'pyodide_entrypoint_helper is required for dedicated snapshot'
+      'customSerializedObjects is required for dedicated snapshot'
     );
   }
-  collectSnapshot(Module, [], pyodide_entrypoint_helper, 'dedicated');
+  collectSnapshot(Module, [], customSerializedObjects, 'dedicated');
 }
 
 /**
@@ -875,7 +887,7 @@ export function maybeCollectDedicatedSnapshot(
  */
 export function maybeCollectSnapshot(
   Module: Module,
-  pyodide_entrypoint_helper: PyodideEntrypointHelper
+  customSerializedObjects: CustomSerializedObjects
 ): void {
   // In order to surface any problems that occur in `memorySnapshotDoImports` to
   // users in local development, always call it even if we aren't actually
@@ -893,21 +905,21 @@ export function maybeCollectSnapshot(
   collectSnapshot(
     Module,
     importedModulesList,
-    pyodide_entrypoint_helper,
+    customSerializedObjects,
     IS_CREATING_BASELINE_SNAPSHOT ? 'baseline' : 'package'
   );
 }
 
 export function finalizeBootstrap(
   Module: Module,
-  pyodide_entrypoint_helper: PyodideEntrypointHelper
+  customSerializedObjects: CustomSerializedObjects
 ): void {
   Module.API.config._makeSnapshot =
     IS_CREATING_SNAPSHOT && Module.API.version !== '0.26.0a2';
   enterJaegerSpan('finalize_bootstrap', () => {
     Module.API.finalizeBootstrap(
       LOADED_SNAPSHOT_META?.hiwire,
-      getHiwireDeserializer(pyodide_entrypoint_helper)
+      getHiwireDeserializer(customSerializedObjects)
     );
   });
   // finalizeBootstrap overrides LD_LIBRARY_PATH. Restore it.

src/pyodide/python-entrypoint-helper.ts

@@ -19,6 +19,7 @@ import {
   WORKFLOWS_ENABLED,
   LEGACY_GLOBAL_HANDLERS,
   LEGACY_INCLUDE_SDK,
+  COMPATIBILITY_FLAGS,
 } from 'pyodide-internal:metadata';
 import { default as Limiter } from 'pyodide-internal:limiter';
 import {
@@ -145,12 +146,10 @@ async function getPyodide(): Promise<Pyodide> {
       return pyodidePromise;
     }
     pyodidePromise = (async function (): Promise<Pyodide> {
-      const pyodide = loadPyodide(
-        IS_WORKERD,
-        LOCKFILE,
-        WORKERD_INDEX_URL,
-        get_pyodide_entrypoint_helper()
-      );
+      const pyodide = loadPyodide(IS_WORKERD, LOCKFILE, WORKERD_INDEX_URL, {
+        pyodide_entrypoint_helper: get_pyodide_entrypoint_helper(),
+        cloudflare_compat_flags: COMPATIBILITY_FLAGS,
+      });
       await setupPatches(pyodide);
       return pyodide;
     })();
@@ -238,6 +237,8 @@ async function setupPatches(pyodide: Pyodide): Promise<void> {
       get_pyodide_entrypoint_helper()
     );
 
+    pyodide.registerJsModule('_cloudflare_compat_flags', COMPATIBILITY_FLAGS);
+
     // Inject modules that enable JS features to be used idiomatically from Python.
     if (LEGACY_INCLUDE_SDK) {
       await injectWorkersApi(pyodide);
@@ -610,10 +611,11 @@ export async function initPython(): Promise<PythonInitResult> {
 
   // Collect a dedicated snapshot at the very end.
   const pyodide = await getPyodide();
-  maybeCollectDedicatedSnapshot(
-    pyodide._module,
-    get_pyodide_entrypoint_helper()
-  );
+  const customSerializedObjects = {
+    pyodide_entrypoint_helper: get_pyodide_entrypoint_helper(),
+    cloudflare_compat_flags: COMPATIBILITY_FLAGS,
+  };
+  maybeCollectDedicatedSnapshot(pyodide._module, customSerializedObjects);
 
   return { handlers, pythonEntrypointClasses, makeEntrypointClass };
 }

src/workerd/server/tests/python/BUILD.bazel

@@ -75,3 +75,5 @@ py_wd_test(
     make_snapshot = False,
     use_snapshot = "numpy",
 )
+
+py_wd_test("python-compat-flag")

src/workerd/server/tests/python/python-compat-flag/python-compat-flag.wd-test

@@ -0,0 +1,14 @@
+using Workerd = import "/workerd/workerd.capnp";
+
+const unitTests :Workerd.Config = (
+  services = [
+    ( name = "python-compat-flag",
+      worker = (
+        modules = [
+          (name = "worker.py", pythonModule = embed "worker.py"),
+        ],
+        compatibilityFlags = [%PYTHON_FEATURE_FLAGS, "python_workflows"],
+      )
+    ),
+  ],
+);

src/workerd/server/tests/python/python-compat-flag/worker.py

@@ -0,0 +1,8 @@
+import _cloudflare_compat_flags
+from workers import WorkerEntrypoint
+
+
+class Default(WorkerEntrypoint):
+    def test(self):
+        assert _cloudflare_compat_flags.python_workflows
+        assert not _cloudflare_compat_flags.python_no_global_handlers