193a74791c8e808ea76001ee07693aa2a5fdd3500johannkoenig@chromium.orgimport unittest
293a74791c8e808ea76001ee07693aa2a5fdd3500johannkoenig@chromium.orgimport textwrap
393a74791c8e808ea76001ee07693aa2a5fdd3500johannkoenig@chromium.orgimport antlr3
493a74791c8e808ea76001ee07693aa2a5fdd3500johannkoenig@chromium.orgimport testbase
593a74791c8e808ea76001ee07693aa2a5fdd3500johannkoenig@chromium.org
693a74791c8e808ea76001ee07693aa2a5fdd3500johannkoenig@chromium.orgclass T(testbase.ANTLRTest):
793a74791c8e808ea76001ee07693aa2a5fdd3500johannkoenig@chromium.org    def testRewrite(self):
893a74791c8e808ea76001ee07693aa2a5fdd3500johannkoenig@chromium.org        self.compileGrammar()
993a74791c8e808ea76001ee07693aa2a5fdd3500johannkoenig@chromium.org
1093a74791c8e808ea76001ee07693aa2a5fdd3500johannkoenig@chromium.org        input = textwrap.dedent(
1193a74791c8e808ea76001ee07693aa2a5fdd3500johannkoenig@chromium.org            '''\
1293a74791c8e808ea76001ee07693aa2a5fdd3500johannkoenig@chromium.org            method foo() {
1393a74791c8e808ea76001ee07693aa2a5fdd3500johannkoenig@chromium.org              i = 3;
1493a74791c8e808ea76001ee07693aa2a5fdd3500johannkoenig@chromium.org              k = i;
1593a74791c8e808ea76001ee07693aa2a5fdd3500johannkoenig@chromium.org              i = k*4;
1693a74791c8e808ea76001ee07693aa2a5fdd3500johannkoenig@chromium.org            }
1793a74791c8e808ea76001ee07693aa2a5fdd3500johannkoenig@chromium.org
1841294d96d7dbf9bc215b09832a8336c5fb158f0bjohannkoenig@chromium.org            method bar() {
1993a74791c8e808ea76001ee07693aa2a5fdd3500johannkoenig@chromium.org              j = i*2;
2093a74791c8e808ea76001ee07693aa2a5fdd3500johannkoenig@chromium.org            }
2193a74791c8e808ea76001ee07693aa2a5fdd3500johannkoenig@chromium.org            ''')
2293a74791c8e808ea76001ee07693aa2a5fdd3500johannkoenig@chromium.org
23ace65784417788374f0b19ce5a8abd06c9ccd007johannkoenig@chromium.org        cStream = antlr3.StringStream(input)
2493a74791c8e808ea76001ee07693aa2a5fdd3500johannkoenig@chromium.org        lexer = self.getLexer(cStream)
2593a74791c8e808ea76001ee07693aa2a5fdd3500johannkoenig@chromium.org        tStream = antlr3.TokenRewriteStream(lexer)
2641294d96d7dbf9bc215b09832a8336c5fb158f0bjohannkoenig@chromium.org        parser = self.getParser(tStream)
2741294d96d7dbf9bc215b09832a8336c5fb158f0bjohannkoenig@chromium.org        parser.program()
2841294d96d7dbf9bc215b09832a8336c5fb158f0bjohannkoenig@chromium.org
2941294d96d7dbf9bc215b09832a8336c5fb158f0bjohannkoenig@chromium.org        expectedOutput = textwrap.dedent('''\
3093a74791c8e808ea76001ee07693aa2a5fdd3500johannkoenig@chromium.org        public class Wrapper {
3193a74791c8e808ea76001ee07693aa2a5fdd3500johannkoenig@chromium.org        public void foo() {
3293a74791c8e808ea76001ee07693aa2a5fdd3500johannkoenig@chromium.org        int k;
3341294d96d7dbf9bc215b09832a8336c5fb158f0bjohannkoenig@chromium.org        int i;
3493a74791c8e808ea76001ee07693aa2a5fdd3500johannkoenig@chromium.org          i = 3;
3593a74791c8e808ea76001ee07693aa2a5fdd3500johannkoenig@chromium.org          k = i;
3693a74791c8e808ea76001ee07693aa2a5fdd3500johannkoenig@chromium.org          i = k*4;
3793a74791c8e808ea76001ee07693aa2a5fdd3500johannkoenig@chromium.org        }
3893a74791c8e808ea76001ee07693aa2a5fdd3500johannkoenig@chromium.org
3941294d96d7dbf9bc215b09832a8336c5fb158f0bjohannkoenig@chromium.org        public void bar() {
4093a74791c8e808ea76001ee07693aa2a5fdd3500johannkoenig@chromium.org        int j;
4193a74791c8e808ea76001ee07693aa2a5fdd3500johannkoenig@chromium.org          j = i*2;
4293a74791c8e808ea76001ee07693aa2a5fdd3500johannkoenig@chromium.org        }
4393a74791c8e808ea76001ee07693aa2a5fdd3500johannkoenig@chromium.org        }
44d95585fb0ec024f6abd96f7b02e0df58019d46afjohannkoenig@chromium.org
45d95585fb0ec024f6abd96f7b02e0df58019d46afjohannkoenig@chromium.org        ''')
46d95585fb0ec024f6abd96f7b02e0df58019d46afjohannkoenig@chromium.org
47d95585fb0ec024f6abd96f7b02e0df58019d46afjohannkoenig@chromium.org        self.failUnlessEqual(
48d95585fb0ec024f6abd96f7b02e0df58019d46afjohannkoenig@chromium.org            str(tStream),
49d95585fb0ec024f6abd96f7b02e0df58019d46afjohannkoenig@chromium.org            expectedOutput
50d95585fb0ec024f6abd96f7b02e0df58019d46afjohannkoenig@chromium.org            )
51d95585fb0ec024f6abd96f7b02e0df58019d46afjohannkoenig@chromium.org
52d95585fb0ec024f6abd96f7b02e0df58019d46afjohannkoenig@chromium.org
53d95585fb0ec024f6abd96f7b02e0df58019d46afjohannkoenig@chromium.orgif __name__ == '__main__':
54d95585fb0ec024f6abd96f7b02e0df58019d46afjohannkoenig@chromium.org    unittest.main()
55d95585fb0ec024f6abd96f7b02e0df58019d46afjohannkoenig@chromium.org
56d95585fb0ec024f6abd96f7b02e0df58019d46afjohannkoenig@chromium.org