From e5ae057cc541d351f59cd09ae578f5dfb7468ff3 Mon Sep 17 00:00:00 2001
From: Dominik Picheta <dominik@cloudflare.com>
Date: Wed, 31 Jan 2024 15:25:56 +0000
Subject: [PATCH] Accept module names with and without a .py extension.

---
 samples/pyodide/config.capnp            |  2 +-
 src/pyodide/python-entrypoint-helper.js | 10 ++++++++--
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/samples/pyodide/config.capnp b/samples/pyodide/config.capnp
index 310c39bcca5..8273ab8ee4b 100644
--- 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
index d9035068c48..a9f79ce6eef 100644
--- 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 {