Files
pytorch/test/cpp/jit/test_class_parser.cpp
Michael Suo dbe850af5b [jit] do the code reorg (#33851)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/33851

Rationale and context described in #33828.

Script to reproduce the move:
https://gist.github.com/suo/16cbefaaeb67ca5a7c6caffd49b7f6e9
ghstack-source-id: 99079645

Test Plan: Make sure CI passes

Reviewed By: jamesr66a

Differential Revision: D20133869

fbshipit-source-id: 390e9241a9c85366d9005c492ac31f10aa96488e
2020-02-27 13:02:51 -08:00

38 lines
1.0 KiB
C++

#include <test/cpp/jit/test_base.h>
#include <torch/csrc/jit/frontend/parser.h>
#include <torch/csrc/jit/frontend/resolver.h>
namespace torch {
namespace jit {
using namespace torch::jit::script;
const auto testSource = R"JIT(
class FooTest:
def __init__(self, x):
self.x = x
def get_x(self):
return self.x
an_attribute : Tensor
)JIT";
void testClassParser() {
Parser p(std::make_shared<Source>(testSource));
std::vector<Def> definitions;
std::vector<Resolver> resolvers;
const auto classDef = ClassDef(p.parseClass());
p.lexer().expect(TK_EOF);
ASSERT_EQ(classDef.name().name(), "FooTest");
ASSERT_EQ(classDef.body().size(), 3);
ASSERT_EQ(Def(classDef.body()[0]).name().name(), "__init__");
ASSERT_EQ(Def(classDef.body()[1]).name().name(), "get_x");
ASSERT_EQ(
Var(Assign(classDef.body()[2]).lhs()).name().name(), "an_attribute");
ASSERT_FALSE(Assign(classDef.body()[2]).rhs().present());
ASSERT_TRUE(Assign(classDef.body()[2]).type().present());
}
} // namespace jit
} // namespace torch