1
// Copyright (c) 2017-2022 Cloudflare, Inc.
2
// Licensed under the Apache 2.0 license found in the LICENSE file or at:
3
//     https://opensource.org/licenses/Apache-2.0
4
 
5
#include "pyodide.h"
6
 
7
#include <kj/test.h>
8
 
9
namespace workerd::api {
10
namespace {
11
 
12
KJ_TEST("getPythonSnapshotRelease") {
13
  capnp::MallocMessageBuilder arena;
14
  // TODO(beta): Factor out FeatureFlags from WorkerBundle.
15
  auto featureFlags = arena.initRoot<CompatibilityFlags>();
16
 
17
  {
18
    auto res = getPythonSnapshotRelease(featureFlags);
19
    KJ_ASSERT(res == kj::none);
20
  }
21
 
22
  featureFlags.setPythonWorkers(true);
23
  {
24
    auto res = KJ_ASSERT_NONNULL(getPythonSnapshotRelease(featureFlags));
25
    KJ_ASSERT(res.getPyodide() == "0.26.0a2");
26
    KJ_ASSERT(res.getFlagName() == "pythonWorkers");
27
  }
28
 
29
  featureFlags.setPythonWorkersDevPyodide(true);
30
  {
31
    auto res = KJ_ASSERT_NONNULL(getPythonSnapshotRelease(featureFlags));
32
    KJ_ASSERT(res.getPyodide() == "dev");
33
    KJ_ASSERT(res.getFlagName() == "pythonWorkersDevPyodide");
34
  }
35
 
36
  featureFlags.setPythonWorkers(false);
37
  {
38
    auto res = KJ_ASSERT_NONNULL(getPythonSnapshotRelease(featureFlags));
39
    KJ_ASSERT(res.getPyodide() == "dev");
40
    KJ_ASSERT(res.getFlagName() == "pythonWorkersDevPyodide");
41
  }
42
 
43
  featureFlags.setPythonWorkers20250116(true);
44
  {
45
    auto res = KJ_ASSERT_NONNULL(getPythonSnapshotRelease(featureFlags));
46
    KJ_ASSERT(res.getPyodide() == "0.28.2");
47
    KJ_ASSERT(res.getFlagName() == "pythonWorkers20250116");
48
  }
49
 
50
  featureFlags.setPythonWorkersDevPyodide(false);
51
  {
52
    auto res = KJ_ASSERT_NONNULL(getPythonSnapshotRelease(featureFlags));
53
    KJ_ASSERT(res.getPyodide() == "0.28.2");
54
    KJ_ASSERT(res.getFlagName() == "pythonWorkers20250116");
55
  }
56
}
57
 
58
KJ_TEST("basic `import` tests") {
59
  auto files = kj::heapArrayBuilder<kj::String>(2);
60
  files.add(kj::str("import a\nimport z"));
61
  files.add(kj::str("import b"));
62
  auto result = pyodide::PythonModuleInfo::parsePythonScriptImports(files.finish());
63
  KJ_REQUIRE(result.size() == 3);
64
  KJ_REQUIRE(result[0] == "a");
65
  KJ_REQUIRE(result[1] == "z");
66
  KJ_REQUIRE(result[2] == "b");
67
}
68
 
69
KJ_TEST("supports whitespace") {
70
  auto files = kj::heapArrayBuilder<kj::String>(1);
71
  files.add(kj::str("import      a\nimport    \n\tz"));
72
  auto result = pyodide::PythonModuleInfo::parsePythonScriptImports(files.finish());
73
  KJ_REQUIRE(result.size() == 2);
74
  KJ_REQUIRE(result[0] == "a");
75
  KJ_REQUIRE(result[1] == "z");
76
}
77
 
78
KJ_TEST("supports windows newlines") {
79
  auto files = kj::heapArrayBuilder<kj::String>(1);
80
  files.add(kj::str("import      a\r\nimport    \r\n\tz"));
81
  auto result = pyodide::PythonModuleInfo::parsePythonScriptImports(files.finish());
82
  KJ_REQUIRE(result.size() == 2);
83
  KJ_REQUIRE(result[0] == "a");
84
  KJ_REQUIRE(result[1] == "z");
85
}
86
 
87
KJ_TEST("basic `from` test") {
88
  auto files = kj::heapArrayBuilder<kj::String>(1);
89
  files.add(kj::str("from x import a,b\nfrom z import y"));
90
  auto result = pyodide::PythonModuleInfo::parsePythonScriptImports(files.finish());
91
  KJ_REQUIRE(result.size() == 2);
92
  KJ_REQUIRE(result[0] == "x");
93
  KJ_REQUIRE(result[1] == "z");
94
}
95
 
96
KJ_TEST("ignores indented blocks") {
97
  auto files = kj::heapArrayBuilder<kj::String>(1);
98
  files.add(kj::str("import a\nif True:\n  import x\nimport y"));
99
  auto result = pyodide::PythonModuleInfo::parsePythonScriptImports(files.finish());
100
  KJ_REQUIRE(result.size() == 2);
101
  KJ_REQUIRE(result[0] == "a");
102
  KJ_REQUIRE(result[1] == "y");
103
}
104
 
105
KJ_TEST("supports nested imports") {
106
  auto files = kj::heapArrayBuilder<kj::String>(1);
107
  files.add(kj::str("import a.b\nimport z.x.y.i"));
108
  auto result = pyodide::PythonModuleInfo::parsePythonScriptImports(files.finish());
109
  KJ_REQUIRE(result.size() == 2);
110
  KJ_REQUIRE(result[0] == "a.b");
111
  KJ_REQUIRE(result[1] == "z.x.y.i");
112
}
113
 
114
KJ_TEST("nested `from` test") {
115
  auto files = kj::heapArrayBuilder<kj::String>(1);
116
  files.add(kj::str("from x.y.z import a,b\nfrom z import y"));
117
  auto result = pyodide::PythonModuleInfo::parsePythonScriptImports(files.finish());
118
  KJ_REQUIRE(result.size() == 2);
119
  KJ_REQUIRE(result[0] == "x.y.z");
120
  KJ_REQUIRE(result[1] == "z");
121
}
122
 
123
KJ_TEST("ignores trailing period") {
124
  auto files = kj::heapArrayBuilder<kj::String>(1);
125
  files.add(kj::str("import a.b.\nimport z.x.y.i."));
126
  auto result = pyodide::PythonModuleInfo::parsePythonScriptImports(files.finish());
127
  KJ_REQUIRE(result.size() == 0);
128
}
129
 
130
KJ_TEST("ignores relative import") {
131
  // This is where we diverge from the old AST-based approach. It would have returned `y` in the
132
  // input below.
133
  auto files = kj::heapArrayBuilder<kj::String>(1);
134
  files.add(kj::str("import .a.b\nimport ..z.x\nfrom .y import x"));
135
  auto result = pyodide::PythonModuleInfo::parsePythonScriptImports(files.finish());
136
  KJ_REQUIRE(result.size() == 0);
137
}
138
 
139
KJ_TEST("supports commas") {
140
  auto files = kj::heapArrayBuilder<kj::String>(1);
141
  files.add(kj::str("import a,b"));
142
  auto result = pyodide::PythonModuleInfo::parsePythonScriptImports(files.finish());
143
  KJ_REQUIRE(result.size() == 2);
144
  KJ_REQUIRE(result[0] == "a");
145
  KJ_REQUIRE(result[1] == "b");
146
}
147
 
148
KJ_TEST("supports backslash") {
149
  auto files = kj::heapArrayBuilder<kj::String>(4);
150
  files.add(kj::str("import a\\\n,b"));
151
  files.add(kj::str("import\\\n q,w"));
152
  files.add(kj::str("from \\\nx import y"));
153
  files.add(kj::str("from \\\n   c import y"));
154
  auto result = pyodide::PythonModuleInfo::parsePythonScriptImports(files.finish());
155
  KJ_REQUIRE(result.size() == 6);
156
  KJ_REQUIRE(result[0] == "a");
157
  KJ_REQUIRE(result[1] == "b");
158
  KJ_REQUIRE(result[2] == "q");
159
  KJ_REQUIRE(result[3] == "w");
160
  KJ_REQUIRE(result[4] == "x");
161
  KJ_REQUIRE(result[5] == "c");
162
}
163
 
164
KJ_TEST("multiline-strings ignored") {
165
  auto files = kj::heapArrayBuilder<kj::String>(4);
166
  files.add(kj::str(R"SCRIPT(
167
FOO="""
168
import x
169
from y import z
170
"""
171
)SCRIPT"));
172
  files.add(kj::str(R"SCRIPT(
173
FOO='''
174
import f
175
from g import z
176
'''
177
)SCRIPT"));
178
  files.add(kj::str(R"SCRIPT(FOO = "\
179
import b \
180
")SCRIPT"));
181
  files.add(kj::str("FOO=\"\"\"  \n", R"SCRIPT(import x
182
from y import z
183
""")SCRIPT"));
184
  auto result = pyodide::PythonModuleInfo::parsePythonScriptImports(files.finish());
185
  KJ_REQUIRE(result.size() == 0);
186
}
187
 
188
KJ_TEST("multiline-strings with imports in-between") {
189
  auto files = kj::heapArrayBuilder<kj::String>(1);
190
  files.add(kj::str(
191
      R"SCRIPT(FOO="""
192
import x
193
from y import z
194
"""import q
195
import w
196
BAR="""
197
import e
198
"""
199
from t import u)SCRIPT"));
200
  auto result = pyodide::PythonModuleInfo::parsePythonScriptImports(files.finish());
201
  KJ_REQUIRE(result.size() == 2);
202
  KJ_REQUIRE(result[0] == "w");
203
  KJ_REQUIRE(result[1] == "t");
204
}
205
 
206
KJ_TEST("import after string literal") {
207
  auto files = kj::heapArrayBuilder<kj::String>(1);
208
  files.add(kj::str(R"SCRIPT(import a
209
"import b)SCRIPT"));
210
  auto result = pyodide::PythonModuleInfo::parsePythonScriptImports(files.finish());
211
  KJ_REQUIRE(result.size() == 1);
212
  KJ_REQUIRE(result[0] == "a");
213
}
214
 
215
KJ_TEST("import after `i`") {
216
  auto files = kj::heapArrayBuilder<kj::String>(1);
217
  files.add(kj::str(R"SCRIPT(import a
218
iimport b)SCRIPT"));
219
  auto result = pyodide::PythonModuleInfo::parsePythonScriptImports(files.finish());
220
  KJ_REQUIRE(result.size() == 1);
221
  KJ_REQUIRE(result[0] == "a");
222
}
223
 
224
KJ_TEST("langchain import") {
225
  auto files = kj::heapArrayBuilder<kj::String>(1);
226
  files.add(kj::str(R"SCRIPT(from js import Response, console, URL
227
from langchain.chat_models import ChatOpenAI
228
import openai)SCRIPT"));
229
  auto result = pyodide::PythonModuleInfo::parsePythonScriptImports(files.finish());
230
  KJ_REQUIRE(result.size() == 3);
231
  KJ_REQUIRE(result[0] == "js");
232
  KJ_REQUIRE(result[1] == "langchain.chat_models");
233
  KJ_REQUIRE(result[2] == "openai");
234
}
235
 
236
KJ_TEST("quote in multiline string") {
237
  auto files = kj::heapArrayBuilder<kj::String>(1);
238
  files.add(kj::str(R"SCRIPT(temp = """
239
w["h
240
""")SCRIPT"));
241
  auto result = pyodide::PythonModuleInfo::parsePythonScriptImports(files.finish());
242
  KJ_REQUIRE(result.size() == 0);
243
}
244
 
245
using pyodide::PythonModuleInfo;
246
 
247
template <typename... Params>
248
kj::Array<kj::String> strArray(Params&&... params) {
249
  return kj::arr(kj::str(params)...);
250
}
251
 
252
template <typename... Params>
253
kj::Array<kj::Array<kj::byte>> bytesArray(Params&&... params) {
254
  return kj::arr(kj::heapArray<kj::byte>(kj::str(params).asBytes())...);
255
}
256
 
257
template <typename... Params>
258
kj::HashSet<kj::String> strSet(Params&&... params) {
259
  auto array = strArray(params...);
260
  kj::HashSet<kj::String> set;
261
  for (auto& str: array) {
262
    set.insert(kj::mv(str));
263
  }
264
  return set;
265
}
266
 
267
KJ_TEST("basic test of getPackageSnapshotImports") {
268
  auto a = pyodide::PythonModuleInfo(strArray("a.py"),
269
      bytesArray("from js import Response\n"
270
                 "import asyncio\n"
271
                 "import numbers\n"
272
                 "def on_fetch(request):\n"
273
                 "  return Response.new('Hello')\n"));
274
  auto result = a.getPackageSnapshotImports("0.26.0a2");
275
  KJ_REQUIRE(result.size() == 1);
276
  KJ_REQUIRE(result[0] == "numbers");
277
};
278
 
279
KJ_TEST("basic test of getPackageSnapshotImports user module") {
280
  auto a = pyodide::PythonModuleInfo(strArray("a.py", "numbers.py"),
281
      bytesArray("from js import Response\n"
282
                 "import asyncio\n"
283
                 "import numbers\n"
284
                 "def on_fetch(request):\n"
285
                 "  return Response.new('Hello')\n",
286
          ""));
287
  auto result = a.getPackageSnapshotImports("0.26.0a2");
288
  KJ_REQUIRE(result.size() == 0);
289
};
290
 
291
kj::Array<kj::String> filterPythonScriptImports(
292
    kj::Array<kj::String> names, kj::ArrayPtr<kj::String> imports, kj::StringPtr version) {
293
  auto contentsBuilder = kj::heapArrayBuilder<kj::Array<kj::byte>>(names.size());
294
  for (auto _: kj::zeroTo(names.size())) {
295
    (void)_;
296
    contentsBuilder.add(kj::Array<kj::byte>(nullptr));
297
  }
298
  auto modInfo = pyodide::PythonModuleInfo(kj::mv(names), contentsBuilder.finish());
299
  auto modSet = modInfo.getWorkerModuleSet();
300
  return PythonModuleInfo::filterPythonScriptImports(kj::mv(modSet), kj::mv(imports), version);
301
}
302
 
303
KJ_TEST("Simple pass through") {
304
  auto imports = strArray("b", "c");
305
  auto result = filterPythonScriptImports({}, kj::mv(imports), "");
306
  KJ_REQUIRE(result.size() == 2);
307
  KJ_REQUIRE(result[0] == "b");
308
  KJ_REQUIRE(result[1] == "c");
309
}
310
 
311
KJ_TEST("pyodide and submodules") {
312
  auto imports = strArray("pyodide", "pyodide.ffi");
313
  auto result = filterPythonScriptImports({}, kj::mv(imports), "0.26.0a2");
314
  KJ_REQUIRE(result.size() == 0);
315
}
316
 
317
KJ_TEST("js and submodules") {
318
  auto imports = strArray("js", "js.crypto");
319
  auto result = filterPythonScriptImports({}, kj::mv(imports), "0.26.0a2");
320
  KJ_REQUIRE(result.size() == 0);
321
}
322
 
323
KJ_TEST("importlib and submodules") {
324
  // importlib and importlib.metadata are imported into the baseline snapshot, but importlib.resources is not.
325
  auto imports = strArray("importlib", "importlib.metadata", "importlib.resources");
326
  auto result = filterPythonScriptImports({}, kj::mv(imports), "");
327
  KJ_REQUIRE(result.size() == 1);
328
  KJ_REQUIRE(result[0] == "importlib.resources");
329
}
330
 
331
KJ_TEST("Filter worker .py files") {
332
  auto workerModules = strArray("b.py", "c.py");
333
  auto imports = strArray("b", "c", "d");
334
  auto result = filterPythonScriptImports(kj::mv(workerModules), kj::mv(imports), "");
335
  KJ_REQUIRE(result.size() == 1);
336
  KJ_REQUIRE(result[0] == "d");
337
}
338
 
339
KJ_TEST("Filter worker module/__init__.py") {
340
  auto workerModules = strArray("a/__init__.py", "b/__init__.py", "c/a.py");
341
  auto imports = strArray("a", "b", "c");
342
  auto result = filterPythonScriptImports(kj::mv(workerModules), kj::mv(imports), "");
343
  KJ_REQUIRE(result.size() == 0);
344
}
345
 
346
KJ_TEST("Filters out subdir/submodule") {
347
  auto workerModules = strArray("subdir/submodule.py");
348
  auto imports = strArray("subdir.submodule");
349
  auto result = filterPythonScriptImports(kj::mv(workerModules), kj::mv(imports), "");
350
  KJ_REQUIRE(result.size() == 0);
351
}
352
 
353
KJ_TEST("Filters out so") {
354
  auto workerModules = strArray("a.so", "b.txt");
355
  auto imports = strArray("a", "b");
356
  auto result = filterPythonScriptImports(kj::mv(workerModules), kj::mv(imports), "");
357
  KJ_REQUIRE(result.size() == 1);
358
  KJ_REQUIRE(result[0] == "b");
359
}
360
 
361
KJ_TEST("Filters out vendor stuff") {
362
  auto workerModules = strArray("python_modules/a.py", "python_modules/package/b.py",
363
      "python_modules/c.so", "python_modules/x.txt");
364
  auto imports = strArray("a", "package", "x");
365
  auto result = filterPythonScriptImports(kj::mv(workerModules), kj::mv(imports), "");
366
  KJ_REQUIRE(result.size() == 1);
367
  KJ_REQUIRE(result[0] == "x");
368
}
369
 
370
}  // namespace
371
}  // namespace workerd::api