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