1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
|
#include <torch/csrc/jit/frontend/tree_views.h>
namespace torch {
namespace jit {
namespace {
void collectUnresolvedNames(
std::vector<std::string>& names,
const TreeView& node) {
if (node.kind() == TK_ASSIGN) {
for (const auto& expr : Assign{node.get()}.lhs_list()) {
collectUnresolvedNames(names, expr);
}
} else if (node.kind() == TK_TUPLE_LITERAL) {
for (const auto& expr : TupleLiteral{node.get()}.inputs()) {
collectUnresolvedNames(names, expr);
}
} else if (node.kind() == TK_LIST_LITERAL) {
for (const auto& expr : ListLiteral{node.get()}.inputs()) {
collectUnresolvedNames(names, expr);
}
} else if (node.kind() == TK_VAR) {
names.push_back(Var{node.get()}.name().name());
}
}
} // namespace
std::vector<std::string> getUnresolvedClassAttributes(const ClassDef& def) {
if (!def.assigns().present()) {
return {};
}
std::vector<std::string> ret;
for (const auto& assign : def.assigns().get()) {
collectUnresolvedNames(ret, assign);
}
return ret;
}
/* static */ ClassDef ClassDef::create(
const SourceRange& range,
const Ident& name,
const Maybe<Expr>& superclass,
const List<Stmt>& body,
const List<Property>& properties,
const List<Assign>& assigns) {
return ClassDef(Compound::create(
TK_CLASS_DEF,
range,
{name,
superclass,
body,
Maybe<List<Property>>::create(range, properties),
Maybe<List<Assign>>::create(range, assigns)}));
}
} // namespace jit
} // namespace torch
|
