1import re
2import unittest
3
4from mako import exceptions
5from mako.codegen import _FOR_LOOP
6from mako.lookup import TemplateLookup
7from mako.runtime import LoopContext
8from mako.runtime import LoopStack
9from mako.template import Template
10from mako.testing.assertions import assert_raises_message
11from mako.testing.fixtures import TemplateTest
12from mako.testing.helpers import flatten_result
13
14
15class TestLoop(unittest.TestCase):
16    def test__FOR_LOOP(self):
17        for statement, target_list, expression_list in (
18            ("for x in y:", "x", "y"),
19            ("for x, y in z:", "x, y", "z"),
20            ("for (x,y) in z:", "(x,y)", "z"),
21            ("for ( x, y, z) in a:", "( x, y, z)", "a"),
22            ("for x in [1, 2, 3]:", "x", "[1, 2, 3]"),
23            ('for x in "spam":', "x", '"spam"'),
24            (
25                "for k,v in dict(a=1,b=2).items():",
26                "k,v",
27                "dict(a=1,b=2).items()",
28            ),
29            (
30                "for x in [y+1 for y in [1, 2, 3]]:",
31                "x",
32                "[y+1 for y in [1, 2, 3]]",
33            ),
34            (
35                "for ((key1, val1), (key2, val2)) in pairwise(dict.items()):",
36                "((key1, val1), (key2, val2))",
37                "pairwise(dict.items())",
38            ),
39            (
40                "for (key1, val1), (key2, val2) in pairwise(dict.items()):",
41                "(key1, val1), (key2, val2)",
42                "pairwise(dict.items())",
43            ),
44        ):
45            match = _FOR_LOOP.match(statement)
46            assert match and match.groups() == (target_list, expression_list)
47
48    def test_no_loop(self):
49        template = Template(
50            """% for x in 'spam':
51${x}
52% endfor"""
53        )
54        code = template.code
55        assert not re.match(r"loop = __M_loop._enter\(:", code), (
56            "No need to "
57            "generate a loop context if the loop variable wasn't accessed"
58        )
59        print(template.render())
60
61    def test_loop_demo(self):
62        template = Template(
63            """x|index|reverse_index|first|last|cycle|even|odd
64% for x in 'ham':
65${x}|${loop.index}|${loop.reverse_index}|${loop.first}|"""
66            """${loop.last}|${loop.cycle('even', 'odd')}|"""
67            """${loop.even}|${loop.odd}
68% endfor"""
69        )
70        expected = [
71            "x|index|reverse_index|first|last|cycle|even|odd",
72            "h|0|2|True|False|even|True|False",
73            "a|1|1|False|False|odd|False|True",
74            "m|2|0|False|True|even|True|False",
75        ]
76        code = template.code
77        assert "loop = __M_loop._enter(" in code, (
78            "Generated a loop context since " "the loop variable was accessed"
79        )
80        rendered = template.render()
81        print(rendered)
82        for line in expected:
83            assert line in rendered, (
84                "Loop variables give information about "
85                "the progress of the loop"
86            )
87
88    def test_nested_loops(self):
89        template = Template(
90            """% for x in 'ab':
91${x} ${loop.index} <- start in outer loop
92% for y in [0, 1]:
93${y} ${loop.index} <- go to inner loop
94% endfor
95${x} ${loop.index} <- back to outer loop
96% endfor"""
97        )
98        rendered = template.render()
99        expected = [
100            "a 0 <- start in outer loop",
101            "0 0 <- go to inner loop",
102            "1 1 <- go to inner loop",
103            "a 0 <- back to outer loop",
104            "b 1 <- start in outer loop",
105            "0 0 <- go to inner loop",
106            "1 1 <- go to inner loop",
107            "b 1 <- back to outer loop",
108        ]
109        for line in expected:
110            assert line in rendered, (
111                "The LoopStack allows you to take "
112                "advantage of the loop variable even in embedded loops"
113            )
114
115    def test_parent_loops(self):
116        template = Template(
117            """% for x in 'ab':
118${x} ${loop.index} <- outer loop
119% for y in [0, 1]:
120${y} ${loop.index} <- inner loop
121${x} ${loop.parent.index} <- parent loop
122% endfor
123${x} ${loop.index} <- outer loop
124% endfor"""
125        )
126        code = template.code
127        rendered = template.render()
128        expected = [
129            "a 0 <- outer loop",
130            "a 0 <- parent loop",
131            "b 1 <- outer loop",
132            "b 1 <- parent loop",
133        ]
134        for line in expected:
135            print(code)
136            assert line in rendered, (
137                "The parent attribute of a loop gives "
138                "you the previous loop context in the stack"
139            )
140
141    def test_out_of_context_access(self):
142        template = Template("""${loop.index}""")
143        assert_raises_message(
144            exceptions.RuntimeException,
145            "No loop context is established",
146            template.render,
147        )
148
149
150class TestLoopStack(unittest.TestCase):
151    def setUp(self):
152        self.stack = LoopStack()
153        self.bottom = "spam"
154        self.stack.stack = [self.bottom]
155
156    def test_enter(self):
157        iterable = "ham"
158        s = self.stack._enter(iterable)
159        assert s is self.stack.stack[-1], (
160            "Calling the stack with an iterable returns " "the stack"
161        )
162        assert iterable == self.stack.stack[-1]._iterable, (
163            "and pushes the " "iterable on the top of the stack"
164        )
165
166    def test__top(self):
167        assert self.bottom == self.stack._top, (
168            "_top returns the last item " "on the stack"
169        )
170
171    def test__pop(self):
172        assert len(self.stack.stack) == 1
173        top = self.stack._pop()
174        assert top == self.bottom
175        assert len(self.stack.stack) == 0
176
177    def test__push(self):
178        assert len(self.stack.stack) == 1
179        iterable = "ham"
180        self.stack._push(iterable)
181        assert len(self.stack.stack) == 2
182        assert iterable is self.stack._top._iterable
183
184    def test_exit(self):
185        iterable = "ham"
186        self.stack._enter(iterable)
187        before = len(self.stack.stack)
188        self.stack._exit()
189        after = len(self.stack.stack)
190        assert before == (after + 1), "Exiting a context pops the stack"
191
192
193class TestLoopContext(unittest.TestCase):
194    def setUp(self):
195        self.iterable = [1, 2, 3]
196        self.ctx = LoopContext(self.iterable)
197
198    def test___len__(self):
199        assert len(self.iterable) == len(self.ctx), (
200            "The LoopContext is the " "same length as the iterable"
201        )
202
203    def test_index(self):
204        expected = tuple(range(len(self.iterable)))
205        actual = tuple(self.ctx.index for i in self.ctx)
206        assert expected == actual, (
207            "The index is consistent with the current " "iteration count"
208        )
209
210    def test_reverse_index(self):
211        length = len(self.iterable)
212        expected = tuple(length - i - 1 for i in range(length))
213        actual = tuple(self.ctx.reverse_index for i in self.ctx)
214        print(expected, actual)
215        assert expected == actual, (
216            "The reverse_index is the number of " "iterations until the end"
217        )
218
219    def test_first(self):
220        expected = (True, False, False)
221        actual = tuple(self.ctx.first for i in self.ctx)
222        assert expected == actual, "first is only true on the first iteration"
223
224    def test_last(self):
225        expected = (False, False, True)
226        actual = tuple(self.ctx.last for i in self.ctx)
227        assert expected == actual, "last is only true on the last iteration"
228
229    def test_even(self):
230        expected = (True, False, True)
231        actual = tuple(self.ctx.even for i in self.ctx)
232        assert expected == actual, "even is true on even iterations"
233
234    def test_odd(self):
235        expected = (False, True, False)
236        actual = tuple(self.ctx.odd for i in self.ctx)
237        assert expected == actual, "odd is true on odd iterations"
238
239    def test_cycle(self):
240        expected = ("a", "b", "a")
241        actual = tuple(self.ctx.cycle("a", "b") for i in self.ctx)
242        assert expected == actual, "cycle endlessly cycles through the values"
243
244
245class TestLoopFlags(TemplateTest):
246    def test_loop_disabled_template(self):
247        self._do_memory_test(
248            """
249            the loop: ${loop}
250        """,
251            "the loop: hi",
252            template_args=dict(loop="hi"),
253            filters=flatten_result,
254            enable_loop=False,
255        )
256
257    def test_loop_disabled_lookup(self):
258        l = TemplateLookup(enable_loop=False)
259        l.put_string(
260            "x",
261            """
262            the loop: ${loop}
263        """,
264        )
265
266        self._do_test(
267            l.get_template("x"),
268            "the loop: hi",
269            template_args=dict(loop="hi"),
270            filters=flatten_result,
271        )
272
273    def test_loop_disabled_override_template(self):
274        self._do_memory_test(
275            """
276            <%page enable_loop="True" />
277            % for i in (1, 2, 3):
278                ${i} ${loop.index}
279            % endfor
280        """,
281            "1 0 2 1 3 2",
282            template_args=dict(loop="hi"),
283            filters=flatten_result,
284            enable_loop=False,
285        )
286
287    def test_loop_disabled_override_lookup(self):
288        l = TemplateLookup(enable_loop=False)
289        l.put_string(
290            "x",
291            """
292            <%page enable_loop="True" />
293            % for i in (1, 2, 3):
294                ${i} ${loop.index}
295            % endfor
296        """,
297        )
298
299        self._do_test(
300            l.get_template("x"),
301            "1 0 2 1 3 2",
302            template_args=dict(loop="hi"),
303            filters=flatten_result,
304        )
305
306    def test_loop_enabled_override_template(self):
307        self._do_memory_test(
308            """
309            <%page enable_loop="True" />
310            % for i in (1, 2, 3):
311                ${i} ${loop.index}
312            % endfor
313        """,
314            "1 0 2 1 3 2",
315            template_args=dict(),
316            filters=flatten_result,
317        )
318
319    def test_loop_enabled_override_lookup(self):
320        l = TemplateLookup()
321        l.put_string(
322            "x",
323            """
324            <%page enable_loop="True" />
325            % for i in (1, 2, 3):
326                ${i} ${loop.index}
327            % endfor
328        """,
329        )
330
331        self._do_test(
332            l.get_template("x"),
333            "1 0 2 1 3 2",
334            template_args=dict(),
335            filters=flatten_result,
336        )
337