Merge pull request #1597 from cloudflare/dominik/python-source-py-ext…

…ension Accept module names with and without a .py extension.
cloudflare · Feb 1, 2024 · 2f37146 · 2f37146
2 parents fffff79 + e5ae057
commit 2f37146
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/samples/pyodide/config.capnp b/samples/pyodide/config.capnp
@@ -22,7 +22,7 @@ const config :Workerd.Config = (
 
 const mainWorker :Workerd.Worker = (
   modules = [
-    (name = "worker", pythonModule = embed "./worker.py"),
+    (name = "worker.py", pythonModule = embed "./worker.py"),
   ],
   compatibilityDate = "2023-12-18",
   compatibilityFlags = ["experimental"],

diff --git a/src/pyodide/python-entrypoint-helper.js b/src/pyodide/python-entrypoint-helper.js
@@ -144,7 +144,9 @@ async function setupPackages(pyodide) {
   const micropipRequirements = [];
   for (const { name, value } of metadata.globals) {
     if (value.pythonModule !== undefined) {
-      pyodide.FS.writeFile(`/session/${name}.py`, value.pythonModule, {
+      // Support both modules names with the `.py` extension as well as without.
+      const pyFilename = name.endsWith(".py") ? name : `${name}.py`;
+      pyodide.FS.writeFile(`/session/${pyFilename}`, value.pythonModule, {
         canOwn: true,
       });
     }
@@ -201,7 +203,11 @@ async function setupPackages(pyodide) {
     );
   }
 
-  return pyodide.pyimport(metadata.mainModule);
+  // The main module can have a `.py` extension, strip it if it exists.
+  const mainName = metadata.mainModule;
+  const mainModule = mainName.endsWith(".py") ? mainName.slice(0, -3) : mainName;
+
+  return pyodide.pyimport(mainModule);
 }
 
 export default {