Python: add support for subdirectory structure (#1635)

src/pyodide/internal/metadatafs.js

@@ -1,32 +1,55 @@
 import { default as MetadataReader } from "pyodide-internal:runtime-generated/metadata";
 import { createReadonlyFS } from "pyodide-internal:readOnlyFS";
 
+function createTree(paths) {
+  const tree = new Map();
+  paths.forEach((elt, idx) => {
+    let subTree = tree;
+    const parts = elt.split("/");
+    const name = parts.pop();
+    for (const part of parts) {
+      let next = subTree.get(part);
+      if (!next) {
+        next = new Map();
+        subTree.set(part, next);
+      }
+      subTree = next;
+    }
+    subTree.set(name, idx);
+  });
+  return tree;
+}
+
 export function createMetadataFS(Module) {
   const TIMESTAMP = Date.now();
   const names = MetadataReader.getNames();
   const sizes = MetadataReader.getSizes();
-  const nameToIndex = new Map(names.map((val, idx) => [val, idx]));
-
+  const rootInfo = createTree(names);
   const FSOps = {
-    getNodeMode(parent, name, index) {
+    getNodeMode(parent, name, info) {
       return {
         permissions: 0o555, // read and execute but not write
-        isDir: index === undefined,
+        isDir: typeof info !== "number",
       };
     },
-    setNodeAttributes(node, index, isDir) {
+    setNodeAttributes(node, info, isDir) {
       node.modtime = TIMESTAMP;
       node.usedBytes = 0;
-      if (!isDir) {
-        node.index = index;
-        node.usedBytes = sizes[index];
+      if (info === undefined) {
+        info = rootInfo;
+      }
+      if (isDir) {
+        node.tree = info;
+      } else {
+        node.index = info;
+        node.usedBytes = sizes[info];
       }
     },
     readdir(node) {
-      return names;
+      return Array.from(node.tree.keys());
     },
     lookup(parent, name) {
-      return nameToIndex.get(name);
+      return parent.tree.get(name);
     },
     read(stream, position, buffer) {
       return MetadataReader.read(stream.node.index, position, buffer);

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

@@ -22,21 +22,35 @@ wd_test(
     ),
 )
 
+# langchain test: disabled for now because it's flaky
+# TODO: reenable this?
+#
+# wd_test(
+#     src = "langchain/langchain.wd-test",
+#     args = ["--experimental"],
+#     data = glob(
+#         [
+#             "langchain/*",
+#         ],
+#         exclude = ["**/*.wd-test"],
+#     ),
+#     # TODO: fails in micropip on windows with:
+#     #
+#     # ValueError: Can't fetch metadata for 'openai'. Please make sure you have entered a correct
+#     # package name and correctly specified index_urls (if you changed them).
+#     target_compatible_with = select({
+#         "@platforms//os:windows": ["@platforms//:incompatible"],
+#         "//conditions:default": [],
+#     }),
+# )
+
 wd_test(
-    src = "langchain/langchain.wd-test",
+    src = "subdirectory/subdirectory.wd-test",
     args = ["--experimental"],
     data = glob(
         [
-            "langchain/*",
+            "subdirectory/**",
         ],
         exclude = ["**/*.wd-test"],
     ),
-    # TODO: fails in micropip on windows with:
-    #
-    # ValueError: Can't fetch metadata for 'openai'. Please make sure you have entered a correct
-    # package name and correctly specified index_urls (if you changed them).
-    target_compatible_with = select({
-        "@platforms//os:windows": ["@platforms//:incompatible"],
-        "//conditions:default": [],
-    }),
 )

src/workerd/server/tests/python/subdirectory/a.py

@@ -0,0 +1,10 @@
+from js import Response
+from subdir.a import x
+
+
+def fetch(request):
+    return Response.new("hello world")
+
+
+def test():
+    print("Hi there, this is a test", x)

src/workerd/server/tests/python/subdirectory/subdir/a.py

@@ -0,0 +1 @@
+x = 7

src/workerd/server/tests/python/subdirectory/subdirectory.wd-test

@@ -0,0 +1,18 @@
+using Workerd = import "/workerd/workerd.capnp";
+
+const unitTests :Workerd.Config = (
+  services = [
+    ( name = "main",
+      worker = (
+        modules = [
+          (name = "a.py", pythonModule = embed "./a.py"),
+          (name = "subdir/__init__.py", pythonModule = embed "./subdir/__init__.py"),
+          (name = "subdir/a.py", pythonModule = embed "./subdir/a.py"),
+        ],
+        compatibilityDate = "2023-12-18",
+        compatibilityFlags = ["experimental"],
+      )
+    ),
+  ],
+);
+