xref: /aosp_15_r20/external/antlr/runtime/Python3/unittests/testtreewizard.py (revision 16467b971bd3e2009fad32dd79016f2c7e421deb)
1 
2 from io import StringIO
3 import os
4 import unittest
5 
6 from antlr3.tree import CommonTreeAdaptor, CommonTree, INVALID_TOKEN_TYPE
7 from antlr3.treewizard import TreeWizard, computeTokenTypes, \
8      TreePatternLexer, EOF, ID, BEGIN, END, PERCENT, COLON, DOT, ARG, \
9      TreePatternParser, \
10      TreePattern, WildcardTreePattern, TreePatternTreeAdaptor
11 
12 
13 class TestComputeTokenTypes(unittest.TestCase):
14     """Test case for the computeTokenTypes function."""
15 
16     def testNone(self):
17         """computeTokenTypes(None) -> {}"""
18 
19         typeMap = computeTokenTypes(None)
20         self.assertIsInstance(typeMap, dict)
21         self.assertEqual(typeMap, {})
22 
23 
24     def testList(self):
25         """computeTokenTypes(['a', 'b']) -> { 'a': 0, 'b': 1 }"""
26 
27         typeMap = computeTokenTypes(['a', 'b'])
28         self.assertIsInstance(typeMap, dict)
29         self.assertEqual(typeMap, { 'a': 0, 'b': 1 })
30 
31 
32 class TestTreePatternLexer(unittest.TestCase):
33     """Test case for the TreePatternLexer class."""
34 
35     def testBegin(self):
36         """TreePatternLexer(): '('"""
37 
38         lexer = TreePatternLexer('(')
39         type = lexer.nextToken()
40         self.assertEqual(type, BEGIN)
41         self.assertEqual(lexer.sval, '')
42         self.assertFalse(lexer.error)
43 
44 
45     def testEnd(self):
46         """TreePatternLexer(): ')'"""
47 
48         lexer = TreePatternLexer(')')
49         type = lexer.nextToken()
50         self.assertEqual(type, END)
51         self.assertEqual(lexer.sval, '')
52         self.assertFalse(lexer.error)
53 
54 
55     def testPercent(self):
56         """TreePatternLexer(): '%'"""
57 
58         lexer = TreePatternLexer('%')
59         type = lexer.nextToken()
60         self.assertEqual(type, PERCENT)
61         self.assertEqual(lexer.sval, '')
62         self.assertFalse(lexer.error)
63 
64 
65     def testDot(self):
66         """TreePatternLexer(): '.'"""
67 
68         lexer = TreePatternLexer('.')
69         type = lexer.nextToken()
70         self.assertEqual(type, DOT)
71         self.assertEqual(lexer.sval, '')
72         self.assertFalse(lexer.error)
73 
74 
75     def testColon(self):
76         """TreePatternLexer(): ':'"""
77 
78         lexer = TreePatternLexer(':')
79         type = lexer.nextToken()
80         self.assertEqual(type, COLON)
81         self.assertEqual(lexer.sval, '')
82         self.assertFalse(lexer.error)
83 
84 
85     def testEOF(self):
86         """TreePatternLexer(): EOF"""
87 
88         lexer = TreePatternLexer('  \n \r \t ')
89         type = lexer.nextToken()
90         self.assertEqual(type, EOF)
91         self.assertEqual(lexer.sval, '')
92         self.assertFalse(lexer.error)
93 
94 
95     def testID(self):
96         """TreePatternLexer(): ID"""
97 
98         lexer = TreePatternLexer('_foo12_bar')
99         type = lexer.nextToken()
100         self.assertEqual(type, ID)
101         self.assertEqual(lexer.sval, '_foo12_bar')
102         self.assertFalse(lexer.error)
103 
104 
105     def testARG(self):
106         """TreePatternLexer(): ARG"""
107 
108         lexer = TreePatternLexer(r'[ \]bla\n]')
109         type = lexer.nextToken()
110         self.assertEqual(type, ARG)
111         self.assertEqual(lexer.sval, r' ]bla\n')
112         self.assertFalse(lexer.error)
113 
114 
115     def testError(self):
116         """TreePatternLexer(): error"""
117 
118         lexer = TreePatternLexer('1')
119         type = lexer.nextToken()
120         self.assertEqual(type, EOF)
121         self.assertEqual(lexer.sval, '')
122         self.assertTrue(lexer.error)
123 
124 
125 class TestTreePatternParser(unittest.TestCase):
126     """Test case for the TreePatternParser class."""
127 
128     def setUp(self):
129         """Setup text fixure
130 
131         We need a tree adaptor, use CommonTreeAdaptor.
132         And a constant list of token names.
133 
134         """
135 
136         self.adaptor = CommonTreeAdaptor()
137         self.tokens = [
138             "", "", "", "", "", "A", "B", "C", "D", "E", "ID", "VAR"
139             ]
140         self.wizard = TreeWizard(self.adaptor, tokenNames=self.tokens)
141 
142 
143     def testSingleNode(self):
144         """TreePatternParser: 'ID'"""
145         lexer = TreePatternLexer('ID')
146         parser = TreePatternParser(lexer, self.wizard, self.adaptor)
147         tree = parser.pattern()
148         self.assertIsInstance(tree, CommonTree)
149         self.assertEqual(tree.getType(), 10)
150         self.assertEqual(tree.getText(), 'ID')
151 
152 
153     def testSingleNodeWithArg(self):
154         """TreePatternParser: 'ID[foo]'"""
155         lexer = TreePatternLexer('ID[foo]')
156         parser = TreePatternParser(lexer, self.wizard, self.adaptor)
157         tree = parser.pattern()
158         self.assertIsInstance(tree, CommonTree)
159         self.assertEqual(tree.getType(), 10)
160         self.assertEqual(tree.getText(), 'foo')
161 
162 
163     def testSingleLevelTree(self):
164         """TreePatternParser: '(A B)'"""
165         lexer = TreePatternLexer('(A B)')
166         parser = TreePatternParser(lexer, self.wizard, self.adaptor)
167         tree = parser.pattern()
168         self.assertIsInstance(tree, CommonTree)
169         self.assertEqual(tree.getType(), 5)
170         self.assertEqual(tree.getText(), 'A')
171         self.assertEqual(tree.getChildCount(), 1)
172         self.assertEqual(tree.getChild(0).getType(), 6)
173         self.assertEqual(tree.getChild(0).getText(), 'B')
174 
175 
176     def testNil(self):
177         """TreePatternParser: 'nil'"""
178         lexer = TreePatternLexer('nil')
179         parser = TreePatternParser(lexer, self.wizard, self.adaptor)
180         tree = parser.pattern()
181         self.assertIsInstance(tree, CommonTree)
182         self.assertEqual(tree.getType(), 0)
183         self.assertIsNone(tree.getText())
184 
185 
186     def testWildcard(self):
187         """TreePatternParser: '(.)'"""
188         lexer = TreePatternLexer('(.)')
189         parser = TreePatternParser(lexer, self.wizard, self.adaptor)
190         tree = parser.pattern()
191         self.assertIsInstance(tree, WildcardTreePattern)
192 
193 
194     def testLabel(self):
195         """TreePatternParser: '(%a:A)'"""
196         lexer = TreePatternLexer('(%a:A)')
197         parser = TreePatternParser(lexer, self.wizard, TreePatternTreeAdaptor())
198         tree = parser.pattern()
199         self.assertIsInstance(tree, TreePattern)
200         self.assertEqual(tree.label, 'a')
201 
202 
203     def testError1(self):
204         """TreePatternParser: ')'"""
205         lexer = TreePatternLexer(')')
206         parser = TreePatternParser(lexer, self.wizard, self.adaptor)
207         tree = parser.pattern()
208         self.assertIsNone(tree)
209 
210 
211     def testError2(self):
212         """TreePatternParser: '()'"""
213         lexer = TreePatternLexer('()')
214         parser = TreePatternParser(lexer, self.wizard, self.adaptor)
215         tree = parser.pattern()
216         self.assertIsNone(tree)
217 
218 
219     def testError3(self):
220         """TreePatternParser: '(A ])'"""
221         lexer = TreePatternLexer('(A ])')
222         parser = TreePatternParser(lexer, self.wizard, self.adaptor)
223         tree = parser.pattern()
224         self.assertIsNone(tree)
225 
226 
227 class TestTreeWizard(unittest.TestCase):
228     """Test case for the TreeWizard class."""
229 
230     def setUp(self):
231         """Setup text fixure
232 
233         We need a tree adaptor, use CommonTreeAdaptor.
234         And a constant list of token names.
235 
236         """
237 
238         self.adaptor = CommonTreeAdaptor()
239         self.tokens = [
240             "", "", "", "", "", "A", "B", "C", "D", "E", "ID", "VAR"
241             ]
242 
243 
244     def testInit(self):
245         """TreeWizard.__init__()"""
246 
247         wiz = TreeWizard(
248             self.adaptor,
249             tokenNames=['a', 'b']
250             )
251 
252         self.assertIs(wiz.adaptor, self.adaptor)
253         self.assertEqual(
254             wiz.tokenNameToTypeMap,
255             { 'a': 0, 'b': 1 }
256             )
257 
258 
259     def testGetTokenType(self):
260         """TreeWizard.getTokenType()"""
261 
262         wiz = TreeWizard(
263             self.adaptor,
264             tokenNames=self.tokens
265             )
266 
267         self.assertEqual(
268             wiz.getTokenType('A'),
269             5
270             )
271 
272         self.assertEqual(
273             wiz.getTokenType('VAR'),
274             11
275             )
276 
277         self.assertEqual(
278             wiz.getTokenType('invalid'),
279             INVALID_TOKEN_TYPE
280             )
281 
282     def testSingleNode(self):
283         wiz = TreeWizard(self.adaptor, self.tokens)
284         t = wiz.create("ID")
285         found = t.toStringTree()
286         expecting = "ID"
287         self.assertEqual(expecting, found)
288 
289 
290     def testSingleNodeWithArg(self):
291         wiz = TreeWizard(self.adaptor, self.tokens)
292         t = wiz.create("ID[foo]")
293         found = t.toStringTree()
294         expecting = "foo"
295         self.assertEqual(expecting, found)
296 
297 
298     def testSingleNodeTree(self):
299         wiz = TreeWizard(self.adaptor, self.tokens)
300         t = wiz.create("(A)")
301         found = t.toStringTree()
302         expecting = "A"
303         self.assertEqual(expecting, found)
304 
305 
306     def testSingleLevelTree(self):
307         wiz = TreeWizard(self.adaptor, self.tokens)
308         t = wiz.create("(A B C D)")
309         found = t.toStringTree()
310         expecting = "(A B C D)"
311         self.assertEqual(expecting, found)
312 
313 
314     def testListTree(self):
315         wiz = TreeWizard(self.adaptor, self.tokens)
316         t = wiz.create("(nil A B C)")
317         found = t.toStringTree()
318         expecting = "A B C"
319         self.assertEqual(expecting, found)
320 
321 
322     def testInvalidListTree(self):
323         wiz = TreeWizard(self.adaptor, self.tokens)
324         t = wiz.create("A B C")
325         self.assertIsNone(t)
326 
327 
328     def testDoubleLevelTree(self):
329         wiz = TreeWizard(self.adaptor, self.tokens)
330         t = wiz.create("(A (B C) (B D) E)")
331         found = t.toStringTree()
332         expecting = "(A (B C) (B D) E)"
333         self.assertEqual(expecting, found)
334 
335 
336     def __simplifyIndexMap(self, indexMap):
337         return dict( # stringify nodes for easy comparing
338             (ttype, [str(node) for node in nodes])
339             for ttype, nodes in indexMap.items()
340             )
341 
342     def testSingleNodeIndex(self):
343         wiz = TreeWizard(self.adaptor, self.tokens)
344         tree = wiz.create("ID")
345         indexMap = wiz.index(tree)
346         found = self.__simplifyIndexMap(indexMap)
347         expecting = { 10: ["ID"] }
348         self.assertEqual(expecting, found)
349 
350 
351     def testNoRepeatsIndex(self):
352         wiz = TreeWizard(self.adaptor, self.tokens)
353         tree = wiz.create("(A B C D)")
354         indexMap = wiz.index(tree)
355         found = self.__simplifyIndexMap(indexMap)
356         expecting = { 8:['D'], 6:['B'], 7:['C'], 5:['A'] }
357         self.assertEqual(expecting, found)
358 
359 
360     def testRepeatsIndex(self):
361         wiz = TreeWizard(self.adaptor, self.tokens)
362         tree = wiz.create("(A B (A C B) B D D)")
363         indexMap = wiz.index(tree)
364         found = self.__simplifyIndexMap(indexMap)
365         expecting = { 8: ['D', 'D'], 6: ['B', 'B', 'B'], 7: ['C'], 5: ['A', 'A'] }
366         self.assertEqual(expecting, found)
367 
368 
369     def testNoRepeatsVisit(self):
370         wiz = TreeWizard(self.adaptor, self.tokens)
371         tree = wiz.create("(A B C D)")
372 
373         elements = []
374         def visitor(node, parent, childIndex, labels):
375             elements.append(str(node))
376 
377         wiz.visit(tree, wiz.getTokenType("B"), visitor)
378 
379         expecting = ['B']
380         self.assertEqual(expecting, elements)
381 
382 
383     def testNoRepeatsVisit2(self):
384         wiz = TreeWizard(self.adaptor, self.tokens)
385         tree = wiz.create("(A B (A C B) B D D)")
386 
387         elements = []
388         def visitor(node, parent, childIndex, labels):
389             elements.append(str(node))
390 
391         wiz.visit(tree, wiz.getTokenType("C"), visitor)
392 
393         expecting = ['C']
394         self.assertEqual(expecting, elements)
395 
396 
397     def testRepeatsVisit(self):
398         wiz = TreeWizard(self.adaptor, self.tokens)
399         tree = wiz.create("(A B (A C B) B D D)")
400 
401         elements = []
402         def visitor(node, parent, childIndex, labels):
403             elements.append(str(node))
404 
405         wiz.visit(tree, wiz.getTokenType("B"), visitor)
406 
407         expecting = ['B', 'B', 'B']
408         self.assertEqual(expecting, elements)
409 
410 
411     def testRepeatsVisit2(self):
412         wiz = TreeWizard(self.adaptor, self.tokens)
413         tree = wiz.create("(A B (A C B) B D D)")
414 
415         elements = []
416         def visitor(node, parent, childIndex, labels):
417             elements.append(str(node))
418 
419         wiz.visit(tree, wiz.getTokenType("A"), visitor)
420 
421         expecting = ['A', 'A']
422         self.assertEqual(expecting, elements)
423 
424 
425     def testRepeatsVisitWithContext(self):
426         wiz = TreeWizard(self.adaptor, self.tokens)
427         tree = wiz.create("(A B (A C B) B D D)")
428 
429         elements = []
430         def visitor(node, parent, childIndex, labels):
431             elements.append('{}@{}[{}]'.format(node, parent, childIndex))
432 
433         wiz.visit(tree, wiz.getTokenType("B"), visitor)
434 
435         expecting = ['B@A[0]', 'B@A[1]', 'B@A[2]']
436         self.assertEqual(expecting, elements)
437 
438 
439     def testRepeatsVisitWithNullParentAndContext(self):
440         wiz = TreeWizard(self.adaptor, self.tokens)
441         tree = wiz.create("(A B (A C B) B D D)")
442 
443         elements = []
444         def visitor(node, parent, childIndex, labels):
445             elements.append(
446                 '{}@{}[{}]'.format(
447                     node, parent or 'nil', childIndex)
448                 )
449 
450         wiz.visit(tree, wiz.getTokenType("A"), visitor)
451 
452         expecting = ['A@nil[0]', 'A@A[1]']
453         self.assertEqual(expecting, elements)
454 
455 
456     def testVisitPattern(self):
457         wiz = TreeWizard(self.adaptor, self.tokens)
458         tree = wiz.create("(A B C (A B) D)")
459 
460         elements = []
461         def visitor(node, parent, childIndex, labels):
462             elements.append(
463                 str(node)
464                 )
465 
466         wiz.visit(tree, '(A B)', visitor)
467 
468         expecting = ['A'] # shouldn't match overall root, just (A B)
469         self.assertEqual(expecting, elements)
470 
471 
472     def testVisitPatternMultiple(self):
473         wiz = TreeWizard(self.adaptor, self.tokens)
474         tree = wiz.create("(A B C (A B) (D (A B)))")
475 
476         elements = []
477         def visitor(node, parent, childIndex, labels):
478             elements.append(
479                 '{}@{}[{}]'.format(node, parent or 'nil', childIndex)
480                 )
481 
482         wiz.visit(tree, '(A B)', visitor)
483 
484         expecting = ['A@A[2]', 'A@D[0]']
485         self.assertEqual(expecting, elements)
486 
487 
488     def testVisitPatternMultipleWithLabels(self):
489         wiz = TreeWizard(self.adaptor, self.tokens)
490         tree = wiz.create("(A B C (A[foo] B[bar]) (D (A[big] B[dog])))")
491 
492         elements = []
493         def visitor(node, parent, childIndex, labels):
494             elements.append(
495                 '{}@{}[{}]{}&{}'.format(
496                     node,
497                     parent or 'nil',
498                     childIndex,
499                     labels['a'],
500                     labels['b'],
501                     )
502                 )
503 
504         wiz.visit(tree, '(%a:A %b:B)', visitor)
505 
506         expecting = ['foo@A[2]foo&bar', 'big@D[0]big&dog']
507         self.assertEqual(expecting, elements)
508 
509 
510     def testParse(self):
511         wiz = TreeWizard(self.adaptor, self.tokens)
512         t = wiz.create("(A B C)")
513         valid = wiz.parse(t, "(A B C)")
514         self.assertTrue(valid)
515 
516 
517     def testParseSingleNode(self):
518         wiz = TreeWizard(self.adaptor, self.tokens)
519         t = wiz.create("A")
520         valid = wiz.parse(t, "A")
521         self.assertTrue(valid)
522 
523 
524     def testParseSingleNodeFails(self):
525         wiz = TreeWizard(self.adaptor, self.tokens)
526         t = wiz.create("A")
527         valid = wiz.parse(t, "B")
528         self.assertFalse(valid)
529 
530 
531     def testParseFlatTree(self):
532         wiz = TreeWizard(self.adaptor, self.tokens)
533         t = wiz.create("(nil A B C)")
534         valid = wiz.parse(t, "(nil A B C)")
535         self.assertTrue(valid)
536 
537 
538     def testParseFlatTreeFails(self):
539         wiz = TreeWizard(self.adaptor, self.tokens)
540         t = wiz.create("(nil A B C)")
541         valid = wiz.parse(t, "(nil A B)")
542         self.assertFalse(valid)
543 
544 
545     def testParseFlatTreeFails2(self):
546         wiz = TreeWizard(self.adaptor, self.tokens)
547         t = wiz.create("(nil A B C)")
548         valid = wiz.parse(t, "(nil A B A)")
549         self.assertFalse(valid)
550 
551 
552     def testWildcard(self):
553         wiz = TreeWizard(self.adaptor, self.tokens)
554         t = wiz.create("(A B C)")
555         valid = wiz.parse(t, "(A . .)")
556         self.assertTrue(valid)
557 
558 
559     def testParseWithText(self):
560         wiz = TreeWizard(self.adaptor, self.tokens)
561         t = wiz.create("(A B[foo] C[bar])")
562         # C pattern has no text arg so despite [bar] in t, no need
563         # to match text--check structure only.
564         valid = wiz.parse(t, "(A B[foo] C)")
565         self.assertTrue(valid)
566 
567 
568     def testParseWithText2(self):
569         wiz = TreeWizard(self.adaptor, self.tokens)
570         t = wiz.create("(A B[T__32] (C (D E[a])))")
571         # C pattern has no text arg so despite [bar] in t, no need
572         # to match text--check structure only.
573         valid = wiz.parse(t, "(A B[foo] C)")
574         self.assertEqual("(A T__32 (C (D a)))", t.toStringTree())
575 
576 
577     def testParseWithTextFails(self):
578         wiz = TreeWizard(self.adaptor, self.tokens)
579         t = wiz.create("(A B C)")
580         valid = wiz.parse(t, "(A[foo] B C)")
581         self.assertFalse(valid) # fails
582 
583 
584     def testParseLabels(self):
585         wiz = TreeWizard(self.adaptor, self.tokens)
586         t = wiz.create("(A B C)")
587         labels = {}
588         valid = wiz.parse(t, "(%a:A %b:B %c:C)", labels)
589         self.assertTrue(valid)
590         self.assertEqual("A", str(labels["a"]))
591         self.assertEqual("B", str(labels["b"]))
592         self.assertEqual("C", str(labels["c"]))
593 
594 
595     def testParseWithWildcardLabels(self):
596         wiz = TreeWizard(self.adaptor, self.tokens)
597         t = wiz.create("(A B C)")
598         labels = {}
599         valid = wiz.parse(t, "(A %b:. %c:.)", labels)
600         self.assertTrue(valid)
601         self.assertEqual("B", str(labels["b"]))
602         self.assertEqual("C", str(labels["c"]))
603 
604 
605     def testParseLabelsAndTestText(self):
606         wiz = TreeWizard(self.adaptor, self.tokens)
607         t = wiz.create("(A B[foo] C)")
608         labels = {}
609         valid = wiz.parse(t, "(%a:A %b:B[foo] %c:C)", labels)
610         self.assertTrue(valid)
611         self.assertEqual("A", str(labels["a"]))
612         self.assertEqual("foo", str(labels["b"]))
613         self.assertEqual("C", str(labels["c"]))
614 
615 
616     def testParseLabelsInNestedTree(self):
617         wiz = TreeWizard(self.adaptor, self.tokens)
618         t = wiz.create("(A (B C) (D E))")
619         labels = {}
620         valid = wiz.parse(t, "(%a:A (%b:B %c:C) (%d:D %e:E) )", labels)
621         self.assertTrue(valid)
622         self.assertEqual("A", str(labels["a"]))
623         self.assertEqual("B", str(labels["b"]))
624         self.assertEqual("C", str(labels["c"]))
625         self.assertEqual("D", str(labels["d"]))
626         self.assertEqual("E", str(labels["e"]))
627 
628 
629     def testEquals(self):
630         wiz = TreeWizard(self.adaptor, self.tokens)
631         t1 = wiz.create("(A B C)")
632         t2 = wiz.create("(A B C)")
633         same = wiz.equals(t1, t2)
634         self.assertTrue(same)
635 
636 
637     def testEqualsWithText(self):
638         wiz = TreeWizard(self.adaptor, self.tokens)
639         t1 = wiz.create("(A B[foo] C)")
640         t2 = wiz.create("(A B[foo] C)")
641         same = wiz.equals(t1, t2)
642         self.assertTrue(same)
643 
644 
645     def testEqualsWithMismatchedText(self):
646         wiz = TreeWizard(self.adaptor, self.tokens)
647         t1 = wiz.create("(A B[foo] C)")
648         t2 = wiz.create("(A B C)")
649         same = wiz.equals(t1, t2)
650         self.assertFalse(same)
651 
652 
653     def testEqualsWithMismatchedList(self):
654         wiz = TreeWizard(self.adaptor, self.tokens)
655         t1 = wiz.create("(A B C)")
656         t2 = wiz.create("(A B A)")
657         same = wiz.equals(t1, t2)
658         self.assertFalse(same)
659 
660 
661     def testEqualsWithMismatchedListLength(self):
662         wiz = TreeWizard(self.adaptor, self.tokens)
663         t1 = wiz.create("(A B C)")
664         t2 = wiz.create("(A B)")
665         same = wiz.equals(t1, t2)
666         self.assertFalse(same)
667 
668 
669     def testFindPattern(self):
670         wiz = TreeWizard(self.adaptor, self.tokens)
671         t = wiz.create("(A B C (A[foo] B[bar]) (D (A[big] B[dog])))")
672         subtrees = wiz.find(t, "(A B)")
673         found = [str(node) for node in subtrees]
674         expecting = ['foo', 'big']
675         self.assertEqual(expecting, found)
676 
677 
678     def testFindTokenType(self):
679         wiz = TreeWizard(self.adaptor, self.tokens)
680         t = wiz.create("(A B C (A[foo] B[bar]) (D (A[big] B[dog])))")
681         subtrees = wiz.find(t, wiz.getTokenType('A'))
682         found = [str(node) for node in subtrees]
683         expecting = ['A', 'foo', 'big']
684         self.assertEqual(expecting, found)
685 
686 
687 
688 if __name__ == "__main__":
689     unittest.main(testRunner=unittest.TextTestRunner(verbosity=2))
690