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 <workerd/rust/python-parser/lib.rs.h>
6
 
7
#include <kj-rs/kj-rs.h>
8
 
9
#include <kj/test.h>
10
 
11
using namespace kj_rs;
12
 
13
using ::edgeworker::rust::python_parser::get_imports;
14
 
15
kj::Array<kj::String> parseImports(kj::ArrayPtr<kj::StringPtr> cpp_modules) {
16
  auto rust_modules = kj::heapArrayBuilder<::rust::Str const>(cpp_modules.size());
17
  for (auto& entry: cpp_modules) {
18
    rust_modules.add(entry.cStr());
19
  }
20
  ::rust::Slice<::rust::Str const> rust_slice(rust_modules.begin(), rust_modules.size());
21
  auto rust_result = get_imports(rust_slice);
22
  return kj::from<RustCopy>(rust_result);
23
}
24
 
25
namespace workerd::api {
26
namespace {
27
 
28
KJ_TEST("basic `import` tests") {
29
  auto result = parseImports(kj::arr("import a\nimport z"_kj, "import b"_kj));
30
  KJ_REQUIRE(result.size() == 3);
31
  KJ_REQUIRE(result[0] == "a");
32
  KJ_REQUIRE(result[1] == "b");
33
  KJ_REQUIRE(result[2] == "z");
34
}
35
 
36
KJ_TEST("supports whitespace") {
37
  auto result = parseImports(kj::arr("import      a\nimport   \\\n\tz"_kj));
38
  KJ_REQUIRE(result.size() == 2);
39
  KJ_REQUIRE(result[0] == "a");
40
  KJ_REQUIRE(result[1] == "z");
41
}
42
 
43
KJ_TEST("supports windows newlines") {
44
  auto result = parseImports(kj::arr("import      a\r\nimport    \\\r\n\tz"_kj));
45
  KJ_REQUIRE(result.size() == 2);
46
  KJ_REQUIRE(result[0] == "a");
47
  KJ_REQUIRE(result[1] == "z");
48
}
49
 
50
KJ_TEST("basic `from` test") {
51
  auto result = parseImports(kj::arr("from x import a,b\nfrom z import y"_kj));
52
  KJ_REQUIRE(result.size() == 2);
53
  KJ_REQUIRE(result[0] == "x");
54
  KJ_REQUIRE(result[1] == "z");
55
}
56
 
57
KJ_TEST("ignores indented blocks") {
58
  auto result = parseImports(kj::arr("import a\nif True:\n  import x\nimport y"_kj));
59
  KJ_REQUIRE(result.size() == 2);
60
  KJ_REQUIRE(result[0] == "a");
61
  KJ_REQUIRE(result[1] == "y");
62
}
63
 
64
KJ_TEST("supports nested imports") {
65
  auto result = parseImports(kj::arr("import a.b\nimport z.x.y.i"_kj));
66
  KJ_REQUIRE(result.size() == 2);
67
  KJ_REQUIRE(result[0] == "a.b");
68
  KJ_REQUIRE(result[1] == "z.x.y.i");
69
}
70
 
71
KJ_TEST("nested `from` test") {
72
  auto result = parseImports(kj::arr("from x.y.z import a,b\nfrom z import y"_kj));
73
  KJ_REQUIRE(result.size() == 2);
74
  KJ_REQUIRE(result[0] == "x.y.z");
75
  KJ_REQUIRE(result[1] == "z");
76
}
77
 
78
KJ_TEST("ignores trailing period") {
79
  auto result = parseImports(kj::arr("import a.b.\nimport z.x.y.i."_kj));
80
  KJ_REQUIRE(result.size() == 0);
81
}
82
 
83
KJ_TEST("ignores relative import") {
84
  // This is where we diverge from the old AST-based approach. It would have returned `y` in the
85
  // input below.
86
  auto result = parseImports(kj::arr("import .a.b\nimport ..z.x\nfrom .y import x"_kj));
87
  KJ_REQUIRE(result.size() == 0);
88
}
89
 
90
KJ_TEST("supports commas") {
91
  auto result = parseImports(kj::arr("import a,b"_kj));
92
  KJ_REQUIRE(result.size() == 2);
93
  KJ_REQUIRE(result[0] == "a");
94
  KJ_REQUIRE(result[1] == "b");
95
}
96
 
97
KJ_TEST("supports backslash") {
98
  // clang-format off
99
  auto result = parseImports(kj::arr(
100
    "import a\\\n,b"_kj,
101
    "import\\\n q,w"_kj,
102
    "from \\\nx import y"_kj,
103
    "from \\\n   c import y"_kj
104
  ));
105
  // clang-format on
106
  KJ_REQUIRE(result.size() == 6);
107
  KJ_REQUIRE(result[0] == "a");
108
  KJ_REQUIRE(result[1] == "b");
109
  KJ_REQUIRE(result[2] == "c");
110
  KJ_REQUIRE(result[3] == "q");
111
  KJ_REQUIRE(result[4] == "w");
112
  KJ_REQUIRE(result[5] == "x");
113
}
114
 
115
KJ_TEST("multiline-strings ignored") {
116
  // clang-format off
117
  auto files = kj::arr(R"SCRIPT(
118
FOO="""
119
import x
120
from y import z
121
"""
122
)SCRIPT"_kj,
123
R"SCRIPT(
124
FOO='''
125
import f
126
from g import z
127
'''
128
)SCRIPT"_kj,
129
R"SCRIPT(FOO = "\
130
import b \
131
")SCRIPT"_kj,
132
"FOO=\"\"\"  \n"_kj,
133
R"SCRIPT(import x
134
from y import z
135
""")SCRIPT"_kj);
136
  // clang-format on
137
  auto result = parseImports(files);
138
  KJ_REQUIRE(result.size() == 0);
139
}
140
 
141
KJ_TEST("multiline-strings with imports in-between") {
142
  auto files = kj::arr(
143
      R"SCRIPT(FOO="""
144
import x
145
from y import z
146
"""
147
import q
148
import w
149
BAR="""
150
import e
151
"""
152
from t import u)SCRIPT"_kj);
153
  auto result = parseImports(files);
154
  KJ_REQUIRE(result.size() == 3);
155
  KJ_REQUIRE(result[0] == "q");
156
  KJ_REQUIRE(result[1] == "t");
157
  KJ_REQUIRE(result[2] == "w");
158
}
159
 
160
KJ_TEST("import after string literal") {
161
  auto files = kj::arr(R"SCRIPT(import a
162
"import b")SCRIPT"_kj);
163
  auto result = parseImports(files);
164
  KJ_REQUIRE(result.size() == 1);
165
  KJ_REQUIRE(result[0] == "a");
166
}
167
 
168
KJ_TEST("langchain import") {
169
  auto files = kj::arr(R"SCRIPT(from js import Response, console, URL
170
from langchain.chat_models import ChatOpenAI
171
import openai)SCRIPT"_kj);
172
  auto result = parseImports(files);
173
  KJ_REQUIRE(result.size() == 3);
174
  KJ_REQUIRE(result[0] == "js");
175
  KJ_REQUIRE(result[1] == "langchain.chat_models");
176
  KJ_REQUIRE(result[2] == "openai");
177
}
178
 
179
}  // namespace
180
}  // namespace workerd::api