1import unittest
2
3from Cython import StringIOTree as stringtree
4
5code = """
6cdef int spam                   # line 1
7
8cdef ham():
9    a = 1
10    b = 2
11    c = 3
12    d = 4
13
14def eggs():
15    pass
16
17cpdef bacon():
18    print spam
19    print 'scotch'
20    print 'tea?'
21    print 'or coffee?'          # line 16
22"""
23
24linemap = dict(enumerate(code.splitlines()))
25
26class TestStringIOTree(unittest.TestCase):
27
28    def setUp(self):
29        self.tree = stringtree.StringIOTree()
30
31    def test_markers(self):
32        assert not self.tree.allmarkers()
33
34    def test_insertion(self):
35        self.write_lines((1, 2, 3))
36        line_4_to_6_insertion_point = self.tree.insertion_point()
37        self.write_lines((7, 8))
38        line_9_to_13_insertion_point = self.tree.insertion_point()
39        self.write_lines((14, 15, 16))
40
41        line_4_insertion_point = line_4_to_6_insertion_point.insertion_point()
42        self.write_lines((5, 6), tree=line_4_to_6_insertion_point)
43
44        line_9_to_12_insertion_point = (
45            line_9_to_13_insertion_point.insertion_point())
46        self.write_line(13, tree=line_9_to_13_insertion_point)
47
48        self.write_line(4, tree=line_4_insertion_point)
49        self.write_line(9, tree=line_9_to_12_insertion_point)
50        line_10_insertion_point = line_9_to_12_insertion_point.insertion_point()
51        self.write_line(11, tree=line_9_to_12_insertion_point)
52        self.write_line(10, tree=line_10_insertion_point)
53        self.write_line(12, tree=line_9_to_12_insertion_point)
54
55        self.assertEqual(self.tree.allmarkers(), range(1, 17))
56        self.assertEqual(code.strip(), self.tree.getvalue().strip())
57
58
59    def write_lines(self, linenos, tree=None):
60        for lineno in linenos:
61            self.write_line(lineno, tree=tree)
62
63    def write_line(self, lineno, tree=None):
64        if tree is None:
65            tree = self.tree
66        tree.markers.append(lineno)
67        tree.write(linemap[lineno] + '\n')
68