Skip to content
rachel/workerd
File

Blob: src/pyodide/python-entrypoint-helper.ts

typescript648 lines
1// Copyright (c) 2026 Cloudflare, Inc.
2// Licensed under the Apache 2.0 license found in the LICENSE file or at:
3//     https://opensource.org/licenses/Apache-2.0
4 
5/* eslint-disable @typescript-eslint/no-unsafe-argument, @typescript-eslint/no-unsafe-assignment, @typescript-eslint/no-unsafe-member-access, @typescript-eslint/no-unsafe-return  */
6// This file is a BUILTIN module that provides the actual implementation for the
7// python-entrypoint.js USER module.
8 
9import { patch_env_helper } from 'pyodide-internal:envHelpers';
10import { enterJaegerSpan } from 'pyodide-internal:jaeger';
11import { default as Limiter } from 'pyodide-internal:limiter';
12import {
13  COMPATIBILITY_FLAGS,
14  IS_WORKERD,
15  LEGACY_GLOBAL_HANDLERS,
16  EXTERNAL_SDK,
17  LOCKFILE,
18  MAIN_MODULE_NAME,
19  SHOULD_SNAPSHOT_TO_DISK,
20  TRANSITIVE_REQUIREMENTS,
21  WORKFLOWS_ENABLED,
22} from 'pyodide-internal:metadata';
23import {
24  beforeRequest,
25  clearSignals,
26  loadPyodide,
27} from 'pyodide-internal:python';
28import { patchLoadPackage } from 'pyodide-internal:setupPackages';
29import {
30  fillSnapshotJsModules,
31  LOADED_SNAPSHOT_TYPE,
32  maybeCollectDedicatedSnapshot,
33} from 'pyodide-internal:snapshot';
34import {
35  PythonUserError,
36  PythonWorkersInternalError,
37  loadPythonMod,
38  reportError,
39} from 'pyodide-internal:util';
40import { PyodideVersion } from 'pyodide-internal:const';
41import { default as introspectionSource } from 'pyodide-internal:introspection.py';
42export { createImportProxy } from 'pyodide-internal:serializeJsModule';
43 
44type PyFuture<T> = Promise<T> & { copy(): PyFuture<T>; destroy(): void };
45 
46const waitUntilPatched = new WeakSet();
47 
48function patchWaitUntil(ctx: {
49  waitUntil: (p: Promise<void> | PyFuture<void>) => void;
50}): void {
51  let tag;
52  try {
53    tag = Object.prototype.toString.call(ctx);
54  } catch (_e) {}
55  if (tag !== '[object ExecutionContext]') {
56    return;
57  }
58  if (waitUntilPatched.has(ctx)) {
59    return;
60  }
61  const origWaitUntil: (p: Promise<void>) => void = ctx.waitUntil.bind(ctx);
62  function waitUntil(p: Promise<void> | PyFuture<void>): void {
63    origWaitUntil(
64      (async function (): Promise<void> {
65        if ('copy' in p) {
66          p = p.copy();
67        }
68        await p;
69        if ('destroy' in p) {
70          p.destroy();
71        }
72      })()
73    );
74  }
75  ctx.waitUntil = waitUntil;
76  waitUntilPatched.add(ctx);
77}
78 
79export type PyodideEntrypointHelper = {
80  doAnImport: (mod: string) => Promise<any>;
81  cloudflareWorkersModule: { env: any };
82  cloudflareSocketsModule: any;
83  workerEntrypoint: any;
84  patchWaitUntil: typeof patchWaitUntil;
85  patch_env_helper: (patch: unknown) => Generator<void>;
86  TEST_ONLY_THROW_PYTHON_WORKERS_INTERNAL_ERROR: (message: string) => never;
87};
88 
89// Function to import JavaScript modules from Python
90let _pyodide_entrypoint_helper: PyodideEntrypointHelper | null = null;
91 
92function get_pyodide_entrypoint_helper(): PyodideEntrypointHelper {
93  if (!_pyodide_entrypoint_helper) {
94    throw new PythonWorkersInternalError(
95      'pyodide_entrypoint_helper is not initialized'
96    );
97  }
98  return _pyodide_entrypoint_helper;
99}
100 
101export async function setDoAnImport(
102  doAnImport: (mod: string) => Promise<any>,
103  workerEntrypoint: any
104): Promise<void> {
105  _pyodide_entrypoint_helper = {
106    doAnImport,
107    cloudflareWorkersModule: await doAnImport('cloudflare:workers'),
108    cloudflareSocketsModule: await doAnImport('cloudflare:sockets'),
109    workerEntrypoint,
110    patchWaitUntil,
111    patch_env_helper,
112    TEST_ONLY_THROW_PYTHON_WORKERS_INTERNAL_ERROR(message: string): never {
113      if (!COMPATIBILITY_FLAGS.experimental) {
114        throw new Error(
115          'TEST_ONLY_THROW_PYTHON_WORKERS_INTERNAL_ERROR requires the experimental compatibility flag'
116        );
117      }
118      throw new PythonWorkersInternalError(message);
119    },
120  };
121  await fillSnapshotJsModules(doAnImport);
122}
123 
124function handleSrcImport(pyodide: Pyodide, e: any): never {
125  // Users may be expecting to import local modules via the `src` directory, which for a default
126  // project structure will fail. This code will add some extra info to the error message to help
127  // them fix it.
128  if (e.name === 'PythonError' && e.type === 'ModuleNotFoundError') {
129    pyodide.runPython(`
130      try:
131        import sys
132        exc = sys.last_value
133        if exc.name == "src":
134          exc.add_note(
135            "If your main module is inside the 'src' directory then your import " +
136            "statement shouldn't include a 'src.' prefix")
137        raise exc
138      finally:
139        del exc
140    `);
141  }
142  throw e;
143}
144 
145async function pyimportMainModule(pyodide: Pyodide): Promise<PyModule> {
146  if (!MAIN_MODULE_NAME.endsWith('.py')) {
147    throw new PythonUserError(
148      'Main module needs to end with a .py file extension'
149    );
150  }
151  const mainModuleName = MAIN_MODULE_NAME.slice(0, -3);
152  if (pyodide.version === PyodideVersion.V0_26_0a2) {
153    return pyodide.pyimport(mainModuleName);
154  } else {
155    return await pyodide._module.API.pyodide_base.pyimport_impl.callPromising(
156      mainModuleName
157    );
158  }
159}
160 
161let pyodidePromise: Promise<Pyodide> | undefined;
162async function getPyodide(): Promise<Pyodide> {
163  return await enterJaegerSpan('get_pyodide', () => {
164    if (pyodidePromise) {
165      return pyodidePromise;
166    }
167    pyodidePromise = (async function (): Promise<Pyodide> {
168      const pyodide = await loadPyodide(IS_WORKERD, LOCKFILE, {
169        pyodide_entrypoint_helper: get_pyodide_entrypoint_helper(),
170        cloudflare_compat_flags: COMPATIBILITY_FLAGS,
171      });
172      await setupPatches(pyodide);
173      return pyodide;
174    })();
175    return pyodidePromise;
176  });
177}
178 
179/**
180 * Import the data from the data module es6 import called jsModName.py into a module called
181 * pyModName.py. The site_packages directory is on the path.
182 */
183async function injectSitePackagesModule(
184  pyodide: Pyodide,
185  jsModName: string,
186  pyModName: string
187): Promise<void> {
188  const mod = await import(`pyodide-internal:${jsModName}.py`);
189  pyodide.FS.writeFile(
190    `${pyodide.FS.sitePackages}/${pyModName}.py`,
191    new Uint8Array(mod.default),
192    { canOwn: true }
193  );
194}
195 
196/**
197 * Put the patch into site_packages and import it.
198 *
199 * TODO: Ideally we should only import the patch lazily when the package that it patches is
200 * imported. Or just apply the patch directly or upstream a fix.
201 */
202async function applyPatch(pyodide: Pyodide, patchName: string): Promise<void> {
203  await injectSitePackagesModule(
204    pyodide,
205    `patches/${patchName}`,
206    patchName + '_patch'
207  );
208  pyodide.pyimport(patchName + '_patch');
209}
210 
211async function injectWorkersApi(pyodide: Pyodide): Promise<void> {
212  if (EXTERNAL_SDK) {
213    pyodide.FS.mkdir(`${pyodide.FS.sitePackages}/workers`);
214    const template = [
215      `err = ModuleNotFoundError("No module named '$MODNAME'", name="$MODNAME")`,
216      `err.add_note("You need to update to workers-py >= 1.90 or to pass disable_python_external_sdk")`,
217      `raise err`,
218    ].join('\n');
219    pyodide.FS.writeFile(
220      `${pyodide.FS.sitePackages}/workers/__init__.py`,
221      template.replaceAll('$MODNAME', 'workers')
222    );
223    pyodide.FS.writeFile(
224      `${pyodide.FS.sitePackages}/asgi.py`,
225      template.replaceAll('$MODNAME', 'asgi')
226    );
227    return;
228  }
229 
230  const sitePackages = pyodide.FS.sitePackages;
231  if (pyodide.version === PyodideVersion.V0_26_0a2) {
232    // Inject at cloudflare.workers for backwards compatibility
233    pyodide.FS.mkdirTree(`${sitePackages}/cloudflare/workers`);
234    await injectSitePackagesModule(
235      pyodide,
236      'workers-api/src/workers/__init__',
237      'cloudflare/workers/__init__'
238    );
239    await injectSitePackagesModule(
240      pyodide,
241      'workers-api/src/workers/_workers',
242      'cloudflare/workers/_workers'
243    );
244  }
245  // The SDK was moved from `cloudflare.workers` to just `workers`.
246  // Create workers package structure with workflows submodule
247  pyodide.FS.mkdir(`${sitePackages}/workers`);
248  await injectSitePackagesModule(
249    pyodide,
250    'workers-api/src/workers/__init__',
251    'workers/__init__'
252  );
253  await injectSitePackagesModule(
254    pyodide,
255    'workers-api/src/workers/_workers',
256    'workers/_workers'
257  );
258  await injectSitePackagesModule(
259    pyodide,
260    'workers-api/src/workers/workflows',
261    'workers/workflows'
262  );
263  await injectSitePackagesModule(pyodide, 'workers-api/src/asgi', 'asgi');
264}
265 
266async function setupPatches(pyodide: Pyodide): Promise<void> {
267  await enterJaegerSpan('setup_patches', async () => {
268    patchLoadPackage(pyodide);
269 
270    // install any extra packages into the site-packages directory
271    // Expose the doAnImport function and global modules to Python globals
272    pyodide.registerJsModule(
273      '_pyodide_entrypoint_helper',
274      get_pyodide_entrypoint_helper()
275    );
276 
277    pyodide.registerJsModule('_cloudflare_compat_flags', COMPATIBILITY_FLAGS);
278 
279    // Inject modules that enable JS features to be used idiomatically from Python.
280    await injectWorkersApi(pyodide);
281 
282    // Install patches as needed
283    if (TRANSITIVE_REQUIREMENTS.has('aiohttp')) {
284      await applyPatch(pyodide, 'aiohttp');
285    }
286    // Other than the oldest version of httpx, we apply the patch at the build step.
287    if (
288      pyodide._module.API.version === PyodideVersion.V0_26_0a2 &&
289      TRANSITIVE_REQUIREMENTS.has('httpx')
290    ) {
291      await applyPatch(pyodide, 'httpx');
292    }
293  });
294}
295 
296let mainModulePromise: Promise<PyModule> | undefined;
297function getMainModule(): Promise<PyModule> {
298  return enterJaegerSpan('get_main_module', async () => {
299    if (mainModulePromise) {
300      return mainModulePromise;
301    }
302    mainModulePromise = (async function (): Promise<PyModule> {
303      const pyodide = await getPyodide();
304      Limiter.beginStartup();
305      try {
306        return await enterJaegerSpan('pyimport_main_module', () =>
307          pyimportMainModule(pyodide)
308        );
309      } catch (e: any) {
310        handleSrcImport(pyodide, e);
311      } finally {
312        Limiter.finishStartup(LOADED_SNAPSHOT_TYPE);
313      }
314    })();
315    return mainModulePromise;
316  });
317}
318 
319async function preparePython(): Promise<PyModule> {
320  try {
321    const pyodide = await getPyodide();
322    const mainModule = await getMainModule();
323    beforeRequest(pyodide._module);
324    return mainModule;
325  } catch (e) {
326    // In edgeworker test suite, without this we get the file name and line number of the exception
327    // but no traceback. This gives us a full traceback.
328    reportError(e as Error);
329  }
330}
331 
332async function doPyCallHelper(
333  relaxed: boolean,
334  pyfunc: PyCallable,
335  args: any[]
336): Promise<any> {
337  const pyodide = await getPyodide();
338  clearSignals(pyodide._module);
339  try {
340    if (pyfunc.callWithOptions) {
341      return await pyfunc.callWithOptions(
342        { relaxed, promising: true },
343        ...args
344      );
345    }
346    if (relaxed) {
347      return await pyfunc.callRelaxed(...args);
348    }
349    return await pyfunc(...args);
350  } catch (e: any) {
351    const pyodide = await getPyodide();
352    handleSrcImport(pyodide, e);
353  }
354}
355 
356function doPyCall(pyfunc: PyCallable, args: any[]): any {
357  return doPyCallHelper(false, pyfunc, args);
358}
359 
360function doRelaxedPyCall(pyfunc: PyCallable, args: any[]): any {
361  return doPyCallHelper(true, pyfunc, args);
362}
363 
364function makeHandler(pyHandlerName: string): Handler {
365  if (
366    pyHandlerName === 'test' &&
367    SHOULD_SNAPSHOT_TO_DISK &&
368    LEGACY_GLOBAL_HANDLERS
369  ) {
370    return async function () {
371      await getPyodide();
372      console.log('Stored snapshot to disk; quitting without running test');
373    };
374  }
375  return async function (...args: any[]) {
376    const mainModule = await enterJaegerSpan(
377      'prep_python',
378      async () => await preparePython()
379    );
380    const handler = mainModule[pyHandlerName];
381    if (!handler) {
382      throw new PythonUserError(
383        `Python entrypoint "${MAIN_MODULE_NAME}" does not export a handler named "${pyHandlerName}"`
384      );
385    }
386    const result = await enterJaegerSpan('python_code', () => {
387      return doRelaxedPyCall(handler, args);
388    });
389 
390    // Support returning a pyodide.ffi.FetchResponse.
391    return result?.js_response ?? result;
392  };
393}
394 
395async function initPyInstance(
396  className: string,
397  args: any[]
398): Promise<PyModule> {
399  const mainModule = await preparePython();
400  const pyClassConstructor = mainModule[className];
401  if (typeof pyClassConstructor !== 'function') {
402    throw new TypeError(
403      `There is no '${className}' class defined in the Python Worker's main module`
404    );
405  }
406  const res = await doPyCall(pyClassConstructor, args);
407  return res as PyModule;
408}
409 
410// https://developers.cloudflare.com/workers/runtime-apis/rpc/reserved-methods/
411const SPECIAL_HANDLER_NAMES = ['fetch', 'connect'];
412const SPECIAL_DO_HANDLER_NAMES = [
413  'alarm',
414  'webSocketMessage',
415  'webSocketClose',
416  'webSocketError',
417];
418 
419function makeEntrypointProxyHandler(
420  pyInstancePromise: Promise<PyModule>,
421  className: string
422): ProxyHandler<any> {
423  return {
424    get(target, prop, receiver): any {
425      if (typeof prop !== 'string') {
426        return Reflect.get(target, prop, receiver);
427      }
428      const isDurableObject = className === 'DurableObject';
429      const isWorkflow = className === 'WorkflowEntrypoint';
430 
431      // Proxy calls to `fetch` to methods named `on_fetch` (and the same for other handlers.)
432      const isKnownHandler = SPECIAL_HANDLER_NAMES.includes(prop);
433      const isKnownDoHandler =
434        isDurableObject && SPECIAL_DO_HANDLER_NAMES.includes(prop);
435      const isFetch = prop === 'fetch';
436      const isWorkflowHandler = isWorkflow && prop === 'run';
437      if ((isKnownHandler || isKnownDoHandler) && LEGACY_GLOBAL_HANDLERS) {
438        prop = 'on_' + prop;
439      }
440 
441      if (
442        !LEGACY_GLOBAL_HANDLERS &&
443        prop === 'test' &&
444        SHOULD_SNAPSHOT_TO_DISK
445      ) {
446        return async function () {
447          await getPyodide();
448          console.log('Stored snapshot to disk; quitting without running test');
449        };
450      }
451 
452      return async function (...args: any[]): Promise<any> {
453        // Check if the requested method exists and if so, call it.
454        const pyInstance = await pyInstancePromise;
455 
456        if (typeof pyInstance[prop] !== 'function') {
457          throw new TypeError(`Method ${prop} does not exist`);
458        }
459 
460        if ((isKnownHandler || isKnownDoHandler) && !isFetch) {
461          return await doPyCallHelper(
462            true,
463            pyInstance[prop] as PyCallable,
464            args
465          );
466        }
467 
468        if (WORKFLOWS_ENABLED && isWorkflowHandler) {
469          // we're hiding this behind a compat flag for now
470          return await doPyCallHelper(
471            true,
472            pyInstance[prop] as PyCallable,
473            args
474          );
475        }
476 
477        const introspectionMod = await getIntrospectionMod();
478 
479        const isRelaxed = isFetch || prop === 'test';
480        return await doPyCall(introspectionMod.wrapper_func, [
481          isRelaxed,
482          pyInstance,
483          prop,
484          ...args,
485        ]);
486      };
487    },
488  };
489}
490 
491function makeEntrypointClass(
492  className: string,
493  classKind: AnyClass,
494  methods: string[]
495): any {
496  const result = class EntrypointWrapper extends classKind {
497    constructor(...args: any[]) {
498      super(...args);
499      // Initialise a Python instance of the class.
500      const pyInstancePromise = initPyInstance(className, args);
501      // We do not know the methods that are defined on the RPC class, so we need a proxy to
502      // support any possible method name.
503      return new Proxy(
504        this,
505        makeEntrypointProxyHandler(pyInstancePromise, classKind.name)
506      );
507    }
508  };
509 
510  // Add dummy functions to the class so that the validator can detect them. These will never get
511  // accessed because of the proxy at runtime.
512  for (let method of methods) {
513    if (
514      SUPPORTED_HANDLER_NAMES.includes(method.slice(3)) &&
515      LEGACY_GLOBAL_HANDLERS
516    ) {
517      // Remove the "on_" prefix.
518      method = method.slice(3);
519    }
520    result.prototype[method] = function (): void {};
521  }
522  return result;
523}
524 
525type IntrospectionMod = {
526  __dict__: PyDict;
527  collect_entrypoint_classes: (mod: PyModule) => PythonEntrypointClasses;
528  wrapper_func: PyCallable;
529};
530 
531let introspectionModPromise: Promise<IntrospectionMod> | null = null;
532async function getIntrospectionMod(): Promise<IntrospectionMod> {
533  if (introspectionModPromise === null) {
534    introspectionModPromise = (async (): Promise<IntrospectionMod> => {
535      const pyodide = await getPyodide();
536      return loadPythonMod(
537        pyodide,
538        'introspection',
539        introspectionSource
540      ) as IntrospectionMod;
541    })();
542  }
543 
544  return introspectionModPromise;
545}
546 
547const SUPPORTED_HANDLER_NAMES = [
548  'fetch',
549  'alarm',
550  'scheduled',
551  'trace',
552  'queue',
553  'pubsub',
554];
555 
556type ExporterClassInfo = {
557  className: string;
558  methodNames: string[];
559};
560 
561type PythonEntrypointClasses = {
562  durableObjects: ExporterClassInfo[];
563  workerEntrypoints: ExporterClassInfo[];
564  workflowEntrypoints: ExporterClassInfo[];
565};
566 
567type PythonInitResult = {
568  handlers: { [handlerName: string]: Handler };
569  pythonEntrypointClasses: PythonEntrypointClasses;
570  makeEntrypointClass: typeof makeEntrypointClass;
571};
572 
573function handleDefaultClass(
574  handlers: PythonInitResult['handlers'],
575  workerEntrypoints: ExporterClassInfo[]
576): void {
577  const index = workerEntrypoints.findIndex(
578    (cls) => cls.className === 'Default'
579  );
580  if (index === -1) {
581    return;
582  }
583  const cls = workerEntrypoints[index]!;
584 
585  // Disallow defining a `Default` WorkerEntrypoint and other "default" top-level handlers.
586  if (Object.keys(handlers).length > 0) {
587    throw new TypeError('Cannot define multiple default entrypoints');
588  }
589 
590  handlers['default'] = makeEntrypointClass(
591    'Default',
592    get_pyodide_entrypoint_helper().workerEntrypoint,
593    cls.methodNames
594  );
595  // Remove the default entrypoint from the list of workerEntrypoints to avoid duplication.
596  workerEntrypoints.splice(index, 1);
597}
598 
599export async function initPython(): Promise<PythonInitResult> {
600  const handlers: {
601    [handlerName: string]: Handler;
602  } = {};
603 
604  let pythonEntrypointClasses: PythonEntrypointClasses = {
605    durableObjects: [],
606    workerEntrypoints: [],
607    workflowEntrypoints: [],
608  };
609 
610  const mainModule = await getMainModule();
611 
612  // In order to get the entrypoint classes exported by the worker, we use a Python module
613  // to introspect the user's main module. So we are effectively using Python to analyse the
614  // classes exported by the user worker here. The class names are then exported from here and
615  // used to create the equivalent JS classes via makeEntrypointClass.
616  const introspectionMod = await getIntrospectionMod();
617  pythonEntrypointClasses =
618    introspectionMod.collect_entrypoint_classes(mainModule);
619  handleDefaultClass(handlers, pythonEntrypointClasses.workerEntrypoints);
620 
621  if (LEGACY_GLOBAL_HANDLERS) {
622    // We add all handlers when running in workerd, so that we can handle the case where the
623    // handler is not defined in our own code and throw a more helpful error. See
624    // undefined-handler.wd-test.
625    const addAllHandlers = IS_WORKERD && !handlers['default'];
626    for (const handlerName of SUPPORTED_HANDLER_NAMES) {
627      const pyHandlerName = 'on_' + handlerName;
628      if (addAllHandlers || typeof mainModule[pyHandlerName] === 'function') {
629        handlers[handlerName] = makeHandler(pyHandlerName);
630      }
631    }
632 
633    if (typeof mainModule.test === 'function') {
634      handlers.test = makeHandler('test');
635    }
636  }
637 
638  // Collect a dedicated snapshot at the very end.
639  const pyodide = await getPyodide();
640  const customSerializedObjects = {
641    pyodide_entrypoint_helper: get_pyodide_entrypoint_helper(),
642    cloudflare_compat_flags: COMPATIBILITY_FLAGS,
643  };
644  maybeCollectDedicatedSnapshot(pyodide._module, customSerializedObjects);
645 
646  return { handlers, pythonEntrypointClasses, makeEntrypointClass };
647}