1 | #include <gtest/gtest.h> |
2 | |
3 | #include <ATen/code_template.h> |
4 | #include <test/cpp/jit/test_utils.h> |
5 | |
6 | namespace torch { |
7 | namespace jit { |
8 | |
9 | static const auto ct = at::jit::CodeTemplate(R"( |
10 | int foo($args) { |
11 | |
12 | $bar |
13 | $bar |
14 | $a+$b |
15 | } |
16 | int commatest(int a${,stuff}) |
17 | int notest(int a${,empty,}) |
18 | )" ); |
19 | static const auto ct_expect = R"( |
20 | int foo(hi, 8) { |
21 | |
22 | what |
23 | on many |
24 | lines... |
25 | 7 |
26 | what |
27 | on many |
28 | lines... |
29 | 7 |
30 | 3+4 |
31 | } |
32 | int commatest(int a, things..., others) |
33 | int notest(int a) |
34 | )" ; |
35 | |
36 | TEST(TestCodeTemplate, Copying) { |
37 | at::jit::TemplateEnv e; |
38 | e.s("hi" , "foo" ); |
39 | e.v("what" , {"is" , "this" }); |
40 | at::jit::TemplateEnv c(e); |
41 | c.s("hi" , "foo2" ); |
42 | ASSERT_EQ(e.s("hi" ), "foo" ); |
43 | ASSERT_EQ(c.s("hi" ), "foo2" ); |
44 | ASSERT_EQ(e.v("what" )[0], "is" ); |
45 | } |
46 | |
47 | TEST(TestCodeTemplate, Formatting) { |
48 | at::jit::TemplateEnv e; |
49 | e.v("args" , {"hi" , "8" }); |
50 | e.v("bar" , {"what\non many\nlines..." , "7" }); |
51 | e.s("a" , "3" ); |
52 | e.s("b" , "4" ); |
53 | e.v("stuff" , {"things..." , "others" }); |
54 | e.v("empty" , {}); |
55 | auto s = ct.format(e); |
56 | // std::cout << "'" << s << "'\n"; |
57 | // std::cout << "'" << ct_expect << "'\n"; |
58 | ASSERT_EQ(s, ct_expect); |
59 | } |
60 | |
61 | } // namespace jit |
62 | } // namespace torch |
63 | |