Merge pull request #2685 from cloudflare/dominik/EW-8677

Implements Python import statement parser to avoid ast module.

src/pyodide/internal/snapshot.ts

@@ -6,13 +6,13 @@ import {
   getSitePackagesPath,
 } from 'pyodide-internal:setupPackages';
 import { default as TarReader } from 'pyodide-internal:packages_tar_reader';
-import processScriptImports from 'pyodide-internal:process_script_imports.py';
 import {
   SHOULD_SNAPSHOT_TO_DISK,
   IS_CREATING_BASELINE_SNAPSHOT,
   MEMORY_SNAPSHOT_READER,
 } from 'pyodide-internal:metadata';
 import { reportError, simpleRunPython } from 'pyodide-internal:util';
+import { default as MetadataReader } from 'pyodide-internal:runtime-generated/metadata';
 
 let LOADED_BASELINE_SNAPSHOT: number;
 
@@ -261,24 +261,28 @@ function memorySnapshotDoImports(Module: Module): Array<string> {
     return [];
   }
 
-  // Process the Python modules in the user worker looking for imports of packages which are not
-  // vendored. Vendored packages are skipped because they may contain sensitive information which
-  // we do not want to include in the package snapshot.
-  //
-  // See process_script_imports.py.
-  const processScriptImportsString = new TextDecoder().decode(
-    new Uint8Array(processScriptImports)
+  // The `importedModules` list will contain all modules that have been imported, including local
+  // modules, the usual `js` and other stdlib modules. We want to filter out local imports, so we
+  // grab them and put them into a set for fast filtering.
+  const localModulePaths: Set<string> = new Set<string>(
+    MetadataReader.getNames()
   );
-  simpleRunPython(Module, processScriptImportsString);
+  const importedModules: Array<string> = ArtifactBundler.constructor
+    // @ts-ignore parsePythonScriptImports is a static method.
+    .parsePythonScriptImports(MetadataReader.getWorkerFiles('py'))
+    .filter((module: string) => {
+      const moduleFilename = module.replace('.', '/') + '.py';
+      return !localModulePaths.has(moduleFilename) && module != 'js';
+    });
 
-  const importedModules: Array<string> = JSON.parse(
-    simpleRunPython(
-      Module,
-      'import sys, json; print(json.dumps(CF_LOADED_MODULES), file=sys.stderr)'
-    )
-  );
+  const deduplicatedModules = [...new Set(importedModules)];
+
+  // Import the modules list so they are included in the snapshot.
+  if (deduplicatedModules.length > 0) {
+    simpleRunPython(Module, 'import ' + deduplicatedModules.join(','));
+  }
 
-  return importedModules;
+  return deduplicatedModules;
 }
 
 function checkLoadedSoFiles(dsoJSON: DylinkInfo): void {

src/pyodide/types/FS.d.ts

@@ -29,6 +29,7 @@ interface FS {
     mode: number
   ): FSNode<Info>;
   isFile: (mode: number) => boolean;
+  readdir: (path: string) => Array<string>;
   genericErrors: Error[];
 }
 

src/pyodide/types/runtime-generated/metadata.d.ts

@@ -7,6 +7,7 @@ declare namespace MetadataReader {
   const getMainModule: () => string;
   const hasMemorySnapshot: () => boolean;
   const getNames: () => string[];
+  const getWorkerFiles: (ext: string) => string[];
   const getSizes: () => number[];
   const readMemorySnapshot: (
     offset: number,

src/workerd/api/BUILD.bazel

@@ -112,6 +112,14 @@ wd_cc_library(
     ],
 )
 
+kj_test(
+    src = "pyodide/pyodide-test.c++",
+    deps = [
+        ":pyodide",
+        "//src/workerd/tests:test-fixture",
+    ],
+)
+
 wd_cc_library(
     name = "data-url",
     srcs = ["data-url.c++"],

src/workerd/api/pyodide/pyodide-test.c++

@@ -0,0 +1,118 @@
+// Copyright (c) 2017-2022 Cloudflare, Inc.
+// Licensed under the Apache 2.0 license found in the LICENSE file or at:
+//     https://opensource.org/licenses/Apache-2.0
+
+#include "pyodide.h"
+#include <kj/test.h>
+
+namespace workerd::api {
+namespace {
+
+KJ_TEST("basic `import` tests") {
+  auto files = kj::heapArrayBuilder<kj::String>(2);
+  files.add(kj::str("import a\nimport z"));
+  files.add(kj::str("import b"));
+  auto result = pyodide::ArtifactBundler::parsePythonScriptImports(files.finish());
+  KJ_REQUIRE(result.size() == 3);
+  KJ_REQUIRE(result[0] == "a");
+  KJ_REQUIRE(result[1] == "z");
+  KJ_REQUIRE(result[2] == "b");
+}
+
+KJ_TEST("supports whitespace") {
+  auto files = kj::heapArrayBuilder<kj::String>(1);
+  files.add(kj::str("import      a\nimport    \n\tz"));
+  auto result = pyodide::ArtifactBundler::parsePythonScriptImports(files.finish());
+  KJ_REQUIRE(result.size() == 2);
+  KJ_REQUIRE(result[0] == "a");
+  KJ_REQUIRE(result[1] == "z");
+}
+
+KJ_TEST("supports windows newlines") {
+  auto files = kj::heapArrayBuilder<kj::String>(1);
+  files.add(kj::str("import      a\r\nimport    \r\n\tz"));
+  auto result = pyodide::ArtifactBundler::parsePythonScriptImports(files.finish());
+  KJ_REQUIRE(result.size() == 2);
+  KJ_REQUIRE(result[0] == "a");
+  KJ_REQUIRE(result[1] == "z");
+}
+
+KJ_TEST("basic `from` test") {
+  auto files = kj::heapArrayBuilder<kj::String>(1);
+  files.add(kj::str("from x import a,b\nfrom z import y"));
+  auto result = pyodide::ArtifactBundler::parsePythonScriptImports(files.finish());
+  KJ_REQUIRE(result.size() == 2);
+  KJ_REQUIRE(result[0] == "x");
+  KJ_REQUIRE(result[1] == "z");
+}
+
+KJ_TEST("ignores indented blocks") {
+  auto files = kj::heapArrayBuilder<kj::String>(1);
+  files.add(kj::str("import a\nif True:\n  import x\nimport y"));
+  auto result = pyodide::ArtifactBundler::parsePythonScriptImports(files.finish());
+  KJ_REQUIRE(result.size() == 2);
+  KJ_REQUIRE(result[0] == "a");
+  KJ_REQUIRE(result[1] == "y");
+}
+
+KJ_TEST("supports nested imports") {
+  auto files = kj::heapArrayBuilder<kj::String>(1);
+  files.add(kj::str("import a.b\nimport z.x.y.i"));
+  auto result = pyodide::ArtifactBundler::parsePythonScriptImports(files.finish());
+  KJ_REQUIRE(result.size() == 2);
+  KJ_REQUIRE(result[0] == "a.b");
+  KJ_REQUIRE(result[1] == "z.x.y.i");
+}
+
+KJ_TEST("nested `from` test") {
+  auto files = kj::heapArrayBuilder<kj::String>(1);
+  files.add(kj::str("from x.y.z import a,b\nfrom z import y"));
+  auto result = pyodide::ArtifactBundler::parsePythonScriptImports(files.finish());
+  KJ_REQUIRE(result.size() == 2);
+  KJ_REQUIRE(result[0] == "x.y.z");
+  KJ_REQUIRE(result[1] == "z");
+}
+
+KJ_TEST("ignores trailing period") {
+  auto files = kj::heapArrayBuilder<kj::String>(1);
+  files.add(kj::str("import a.b.\nimport z.x.y.i."));
+  auto result = pyodide::ArtifactBundler::parsePythonScriptImports(files.finish());
+  KJ_REQUIRE(result.size() == 0);
+}
+
+KJ_TEST("ignores relative import") {
+  // This is where we diverge from the old AST-based approach. It would have returned `y` in the
+  // input below.
+  auto files = kj::heapArrayBuilder<kj::String>(1);
+  files.add(kj::str("import .a.b\nimport ..z.x\nfrom .y import x"));
+  auto result = pyodide::ArtifactBundler::parsePythonScriptImports(files.finish());
+  KJ_REQUIRE(result.size() == 0);
+}
+
+KJ_TEST("supports commas") {
+  auto files = kj::heapArrayBuilder<kj::String>(1);
+  files.add(kj::str("import a,b"));
+  auto result = pyodide::ArtifactBundler::parsePythonScriptImports(files.finish());
+  KJ_REQUIRE(result.size() == 2);
+  KJ_REQUIRE(result[0] == "a");
+  KJ_REQUIRE(result[1] == "b");
+}
+
+KJ_TEST("supports backslash") {
+  auto files = kj::heapArrayBuilder<kj::String>(4);
+  files.add(kj::str("import a\\\n,b"));
+  files.add(kj::str("import\\\n q,w"));
+  files.add(kj::str("from \\\nx import y"));
+  files.add(kj::str("from \\\n   c import y"));
+  auto result = pyodide::ArtifactBundler::parsePythonScriptImports(files.finish());
+  KJ_REQUIRE(result.size() == 6);
+  KJ_REQUIRE(result[0] == "a");
+  KJ_REQUIRE(result[1] == "b");
+  KJ_REQUIRE(result[2] == "q");
+  KJ_REQUIRE(result[3] == "w");
+  KJ_REQUIRE(result[4] == "x");
+  KJ_REQUIRE(result[5] == "c");
+}
+
+}  // namespace
+}  // namespace workerd::api