Fixes parsing of imports in multi-line string literals.

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

@@ -115,5 +115,77 @@ KJ_TEST("supports backslash") {
   KJ_REQUIRE(result[5] == "c");
 }
 
+KJ_TEST("multiline-strings ignored") {
+  auto files = kj::heapArrayBuilder<kj::String>(4);
+  files.add(kj::str(R"SCRIPT(
+FOO="""
+import x
+from y import z
+"""
+)SCRIPT"));
+  files.add(kj::str(R"SCRIPT(
+FOO='''
+import f
+from g import z
+'''
+)SCRIPT"));
+  files.add(kj::str(R"SCRIPT(FOO = "\
+import b \
+")SCRIPT"));
+  files.add(kj::str("FOO=\"\"\"  \n", R"SCRIPT(import x
+from y import z
+""")SCRIPT"));
+  auto result = pyodide::ArtifactBundler::parsePythonScriptImports(files.finish());
+  KJ_REQUIRE(result.size() == 0);
+}
+
+KJ_TEST("multiline-strings with imports in-between") {
+  auto files = kj::heapArrayBuilder<kj::String>(1);
+  files.add(kj::str(
+      R"SCRIPT(FOO="""
+import x
+from y import z
+"""import q
+import w
+BAR="""
+import e
+"""
+from t import u)SCRIPT"));
+  auto result = pyodide::ArtifactBundler::parsePythonScriptImports(files.finish());
+  KJ_REQUIRE(result.size() == 2);
+  KJ_REQUIRE(result[0] == "w");
+  KJ_REQUIRE(result[1] == "t");
+}
+
+KJ_TEST("import after string literal") {
+  auto files = kj::heapArrayBuilder<kj::String>(1);
+  files.add(kj::str(R"SCRIPT(import a
+"import b)SCRIPT"));
+  auto result = pyodide::ArtifactBundler::parsePythonScriptImports(files.finish());
+  KJ_REQUIRE(result.size() == 1);
+  KJ_REQUIRE(result[0] == "a");
+}
+
+KJ_TEST("import after `i`") {
+  auto files = kj::heapArrayBuilder<kj::String>(1);
+  files.add(kj::str(R"SCRIPT(import a
+iimport b)SCRIPT"));
+  auto result = pyodide::ArtifactBundler::parsePythonScriptImports(files.finish());
+  KJ_REQUIRE(result.size() == 1);
+  KJ_REQUIRE(result[0] == "a");
+}
+
+KJ_TEST("langchain import") {
+  auto files = kj::heapArrayBuilder<kj::String>(1);
+  files.add(kj::str(R"SCRIPT(from js import Response, console, URL
+from langchain.chat_models import ChatOpenAI
+import openai)SCRIPT"));
+  auto result = pyodide::ArtifactBundler::parsePythonScriptImports(files.finish());
+  KJ_REQUIRE(result.size() == 3);
+  KJ_REQUIRE(result[0] == "js");
+  KJ_REQUIRE(result[1] == "langchain.chat_models");
+  KJ_REQUIRE(result[2] == "openai");
+}
+
 }  // namespace
 }  // namespace workerd::api

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

@@ -207,12 +207,14 @@ kj::Array<kj::String> ArtifactBundler::parsePythonScriptImports(kj::Array<kj::St
 
     int i = 0;
     while (i < file.size()) {
-      auto keywordToParse = file[i] == 'i' ? "import"_kj : "from"_kj;
       switch (file[i]) {
         case 'i':
-        case 'f':
+        case 'f': {
+          auto keywordToParse = file[i] == 'i' ? "import"_kj : "from"_kj;
           if (!parseKeyword(file, keywordToParse, i)) {
-            i += skipUntil(file, {'\n', '\r'}, i);
+            // We cannot simply skip the current char here, doing so would mean that
+            // `iimport x` would be parsed as a valid import.
+            i += skipUntil(file, {'\n', '\r', '"', '\''}, i);
             continue;
           }
           i += keywordToParse.size();  // skip "import" or "from"
@@ -256,9 +258,41 @@ kj::Array<kj::String> ArtifactBundler::parsePythonScriptImports(kj::Array<kj::St
             }
           }
           break;
+        }
+        case '"':
+        case '\'': {
+          char quote = file[i];
+          // Detect multi-line string literals `"""` and skip until the corresponding ending `"""`.
+          if (i + 2 < file.size() && file[i + 1] == quote && file[i + 2] == quote) {
+            i += 3;  // skip start quotes.
+            // skip until terminating quotes.
+            while (i + 2 < file.size() && file[i + 1] != quote && file[i + 2] != quote) {
+              i += skipUntil(file, {quote}, i);
+            }
+            i += 3;  // skip terminating quotes.
+          } else if (i + 2 < file.size() && file[i + 1] == '\\' &&
+              (file[i + 2] == '\n' || file[i + 2] == '\r')) {
+            // Detect string literal with backslash.
+            i += 3;  // skip `"\<NL>`
+            // skip until quote, but ignore `\"`.
+            while (file[i] != quote && file[i - 1] != '\\') {
+              i += skipUntil(file, {quote}, i);
+            }
+            i += 1;  // skip quote.
+          } else {
+            i += 1;  // skip quote.
+          }
+
+          // skip until EOL so that we don't mistakenly parse and capture `"import x`.
+          i += skipUntil(file, {'\n', '\r', '"', '\''}, i);
+          break;
+        }
         default:
-          // Skip to the next line.
-          i += skipUntil(file, {'\n', '\r'}, i);
+          // Skip to the next line or " or '
+          i += skipUntil(file, {'\n', '\r', '"', '\''}, i);
+          if (file[i] == '"' || file[i] == '\'') {
+            continue;  // Allow the quotes to be handled above.
+          }
           if (file[i] != '\0') {
             i += skipChar(file, {'\n', '\r'}, i);  // skip newline.
           }