• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1import doctest
2import traceback
3import unittest
4
5from test.support import BrokenIter
6
7
8doctests = """
9########### Tests mostly copied from test_listcomps.py ############
10
11Test simple loop with conditional
12
13    >>> sum({i*i for i in range(100) if i&1 == 1})
14    166650
15
16Test simple case
17
18    >>> {2*y + x + 1 for x in (0,) for y in (1,)}
19    {3}
20
21Test simple nesting
22
23    >>> list(sorted({(i,j) for i in range(3) for j in range(4)}))
24    [(0, 0), (0, 1), (0, 2), (0, 3), (1, 0), (1, 1), (1, 2), (1, 3), (2, 0), (2, 1), (2, 2), (2, 3)]
25
26Test nesting with the inner expression dependent on the outer
27
28    >>> list(sorted({(i,j) for i in range(4) for j in range(i)}))
29    [(1, 0), (2, 0), (2, 1), (3, 0), (3, 1), (3, 2)]
30
31Test the idiom for temporary variable assignment in comprehensions.
32
33    >>> sorted({j*j for i in range(4) for j in [i+1]})
34    [1, 4, 9, 16]
35    >>> sorted({j*k for i in range(4) for j in [i+1] for k in [j+1]})
36    [2, 6, 12, 20]
37    >>> sorted({j*k for i in range(4) for j, k in [(i+1, i+2)]})
38    [2, 6, 12, 20]
39
40Not assignment
41
42    >>> sorted({i*i for i in [*range(4)]})
43    [0, 1, 4, 9]
44    >>> sorted({i*i for i in (*range(4),)})
45    [0, 1, 4, 9]
46
47Make sure the induction variable is not exposed
48
49    >>> i = 20
50    >>> sum({i*i for i in range(100)})
51    328350
52
53    >>> i
54    20
55
56Verify that syntax error's are raised for setcomps used as lvalues
57
58    >>> {y for y in (1,2)} = 10          # doctest: +IGNORE_EXCEPTION_DETAIL
59    Traceback (most recent call last):
60       ...
61    SyntaxError: ...
62
63    >>> {y for y in (1,2)} += 10         # doctest: +IGNORE_EXCEPTION_DETAIL
64    Traceback (most recent call last):
65       ...
66    SyntaxError: ...
67
68
69Make a nested set comprehension that acts like set(range())
70
71    >>> def srange(n):
72    ...     return {i for i in range(n)}
73    >>> list(sorted(srange(10)))
74    [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
75
76Same again, only as a lambda expression instead of a function definition
77
78    >>> lrange = lambda n:  {i for i in range(n)}
79    >>> list(sorted(lrange(10)))
80    [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
81
82Generators can call other generators:
83
84    >>> def grange(n):
85    ...     for x in {i for i in range(n)}:
86    ...         yield x
87    >>> list(sorted(grange(5)))
88    [0, 1, 2, 3, 4]
89
90
91Make sure that None is a valid return value
92
93    >>> {None for i in range(10)}
94    {None}
95
96########### Tests for various scoping corner cases ############
97
98Return lambdas that use the iteration variable as a default argument
99
100    >>> items = {(lambda i=i: i) for i in range(5)}
101    >>> {x() for x in items} == set(range(5))
102    True
103
104Same again, only this time as a closure variable
105
106    >>> items = {(lambda: i) for i in range(5)}
107    >>> {x() for x in items}
108    {4}
109
110Another way to test that the iteration variable is local to the list comp
111
112    >>> items = {(lambda: i) for i in range(5)}
113    >>> i = 20
114    >>> {x() for x in items}
115    {4}
116
117And confirm that a closure can jump over the list comp scope
118
119    >>> items = {(lambda: y) for i in range(5)}
120    >>> y = 2
121    >>> {x() for x in items}
122    {2}
123
124We also repeat each of the above scoping tests inside a function
125
126    >>> def test_func():
127    ...     items = {(lambda i=i: i) for i in range(5)}
128    ...     return {x() for x in items}
129    >>> test_func() == set(range(5))
130    True
131
132    >>> def test_func():
133    ...     items = {(lambda: i) for i in range(5)}
134    ...     return {x() for x in items}
135    >>> test_func()
136    {4}
137
138    >>> def test_func():
139    ...     items = {(lambda: i) for i in range(5)}
140    ...     i = 20
141    ...     return {x() for x in items}
142    >>> test_func()
143    {4}
144
145    >>> def test_func():
146    ...     items = {(lambda: y) for i in range(5)}
147    ...     y = 2
148    ...     return {x() for x in items}
149    >>> test_func()
150    {2}
151
152"""
153
154class SetComprehensionTest(unittest.TestCase):
155    def test_exception_locations(self):
156        # The location of an exception raised from __init__ or
157        # __next__ should should be the iterator expression
158
159        def init_raises():
160            try:
161                {x for x in BrokenIter(init_raises=True)}
162            except Exception as e:
163                return e
164
165        def next_raises():
166            try:
167                {x for x in BrokenIter(next_raises=True)}
168            except Exception as e:
169                return e
170
171        def iter_raises():
172            try:
173                {x for x in BrokenIter(iter_raises=True)}
174            except Exception as e:
175                return e
176
177        for func, expected in [(init_raises, "BrokenIter(init_raises=True)"),
178                               (next_raises, "BrokenIter(next_raises=True)"),
179                               (iter_raises, "BrokenIter(iter_raises=True)"),
180                              ]:
181            with self.subTest(func):
182                exc = func()
183                f = traceback.extract_tb(exc.__traceback__)[0]
184                indent = 16
185                co = func.__code__
186                self.assertEqual(f.lineno, co.co_firstlineno + 2)
187                self.assertEqual(f.end_lineno, co.co_firstlineno + 2)
188                self.assertEqual(f.line[f.colno - indent : f.end_colno - indent],
189                                 expected)
190
191__test__ = {'doctests' : doctests}
192
193def load_tests(loader, tests, pattern):
194    tests.addTest(doctest.DocTestSuite())
195    return tests
196
197
198if __name__ == "__main__":
199    unittest.main()
200