blob: 85e09a1f575c336cfc774e66b8ff158c39655946 [file] [log] [blame]
Antoine Pitrou796564c2013-07-30 19:59:21 +02001import gc
2import sys
3import unittest
4import weakref
5
6from test import support
7
8
9class FinalizationTest(unittest.TestCase):
10
11 def test_frame_resurrect(self):
12 # A generator frame can be resurrected by a generator's finalization.
13 def gen():
14 nonlocal frame
15 try:
16 yield
17 finally:
18 frame = sys._getframe()
19
20 g = gen()
21 wr = weakref.ref(g)
22 next(g)
23 del g
24 support.gc_collect()
25 self.assertIs(wr(), None)
26 self.assertTrue(frame)
27 del frame
28 support.gc_collect()
29
30 def test_refcycle(self):
31 # A generator caught in a refcycle gets finalized anyway.
32 old_garbage = gc.garbage[:]
33 finalized = False
34 def gen():
35 nonlocal finalized
36 try:
37 g = yield
38 yield 1
39 finally:
40 finalized = True
41
42 g = gen()
43 next(g)
44 g.send(g)
45 self.assertGreater(sys.getrefcount(g), 2)
46 self.assertFalse(finalized)
47 del g
48 support.gc_collect()
49 self.assertTrue(finalized)
50 self.assertEqual(gc.garbage, old_garbage)
51
Serhiy Storchakac775ad62015-03-11 18:20:35 +020052 def test_lambda_generator(self):
53 # Issue #23192: Test that a lambda returning a generator behaves
54 # like the equivalent function
55 f = lambda: (yield 1)
56 def g(): return (yield 1)
57
58 # test 'yield from'
59 f2 = lambda: (yield from g())
60 def g2(): return (yield from g())
61
62 f3 = lambda: (yield from f())
63 def g3(): return (yield from f())
64
65 for gen_fun in (f, g, f2, g2, f3, g3):
66 gen = gen_fun()
67 self.assertEqual(next(gen), 1)
68 with self.assertRaises(StopIteration) as cm:
69 gen.send(2)
70 self.assertEqual(cm.exception.value, 2)
71
Antoine Pitrou796564c2013-07-30 19:59:21 +020072
Victor Stinner40ee3012014-06-16 15:59:28 +020073class GeneratorTest(unittest.TestCase):
74
75 def test_name(self):
76 def func():
77 yield 1
78
79 # check generator names
80 gen = func()
81 self.assertEqual(gen.__name__, "func")
82 self.assertEqual(gen.__qualname__,
83 "GeneratorTest.test_name.<locals>.func")
84
85 # modify generator names
86 gen.__name__ = "name"
87 gen.__qualname__ = "qualname"
88 self.assertEqual(gen.__name__, "name")
89 self.assertEqual(gen.__qualname__, "qualname")
90
91 # generator names must be a string and cannot be deleted
92 self.assertRaises(TypeError, setattr, gen, '__name__', 123)
93 self.assertRaises(TypeError, setattr, gen, '__qualname__', 123)
94 self.assertRaises(TypeError, delattr, gen, '__name__')
95 self.assertRaises(TypeError, delattr, gen, '__qualname__')
96
97 # modify names of the function creating the generator
98 func.__qualname__ = "func_qualname"
99 func.__name__ = "func_name"
100 gen = func()
101 self.assertEqual(gen.__name__, "func_name")
102 self.assertEqual(gen.__qualname__, "func_qualname")
103
104 # unnamed generator
105 gen = (x for x in range(10))
106 self.assertEqual(gen.__name__,
107 "<genexpr>")
108 self.assertEqual(gen.__qualname__,
109 "GeneratorTest.test_name.<locals>.<genexpr>")
110
111
Victor Stinner26f7b8a2015-01-31 10:29:47 +0100112class ExceptionTest(unittest.TestCase):
113 # Tests for the issue #23353: check that the currently handled exception
114 # is correctly saved/restored in PyEval_EvalFrameEx().
115
116 def test_except_throw(self):
117 def store_raise_exc_generator():
118 try:
119 self.assertEqual(sys.exc_info()[0], None)
120 yield
121 except Exception as exc:
122 # exception raised by gen.throw(exc)
123 self.assertEqual(sys.exc_info()[0], ValueError)
124 self.assertIsNone(exc.__context__)
125 yield
126
127 # ensure that the exception is not lost
128 self.assertEqual(sys.exc_info()[0], ValueError)
129 yield
130
131 # we should be able to raise back the ValueError
132 raise
133
134 make = store_raise_exc_generator()
135 next(make)
136
137 try:
138 raise ValueError()
139 except Exception as exc:
140 try:
141 make.throw(exc)
142 except Exception:
143 pass
144
145 next(make)
146 with self.assertRaises(ValueError) as cm:
147 next(make)
148 self.assertIsNone(cm.exception.__context__)
149
150 self.assertEqual(sys.exc_info(), (None, None, None))
151
152 def test_except_next(self):
153 def gen():
154 self.assertEqual(sys.exc_info()[0], ValueError)
155 yield "done"
156
157 g = gen()
158 try:
159 raise ValueError
160 except Exception:
161 self.assertEqual(next(g), "done")
162 self.assertEqual(sys.exc_info(), (None, None, None))
163
164 def test_except_gen_except(self):
165 def gen():
166 try:
167 self.assertEqual(sys.exc_info()[0], None)
168 yield
169 # we are called from "except ValueError:", TypeError must
170 # inherit ValueError in its context
171 raise TypeError()
172 except TypeError as exc:
173 self.assertEqual(sys.exc_info()[0], TypeError)
174 self.assertEqual(type(exc.__context__), ValueError)
175 # here we are still called from the "except ValueError:"
176 self.assertEqual(sys.exc_info()[0], ValueError)
177 yield
178 self.assertIsNone(sys.exc_info()[0])
179 yield "done"
180
181 g = gen()
182 next(g)
183 try:
184 raise ValueError
185 except Exception:
186 next(g)
187
188 self.assertEqual(next(g), "done")
189 self.assertEqual(sys.exc_info(), (None, None, None))
190
191 def test_except_throw_exception_context(self):
192 def gen():
193 try:
194 try:
195 self.assertEqual(sys.exc_info()[0], None)
196 yield
197 except ValueError:
198 # we are called from "except ValueError:"
199 self.assertEqual(sys.exc_info()[0], ValueError)
200 raise TypeError()
201 except Exception as exc:
202 self.assertEqual(sys.exc_info()[0], TypeError)
203 self.assertEqual(type(exc.__context__), ValueError)
204 # we are still called from "except ValueError:"
205 self.assertEqual(sys.exc_info()[0], ValueError)
206 yield
207 self.assertIsNone(sys.exc_info()[0])
208 yield "done"
209
210 g = gen()
211 next(g)
212 try:
213 raise ValueError
214 except Exception as exc:
215 g.throw(exc)
216
217 self.assertEqual(next(g), "done")
218 self.assertEqual(sys.exc_info(), (None, None, None))
219
220
Tim Peters6ba5f792001-06-23 20:45:43 +0000221tutorial_tests = """
Tim Peters1def3512001-06-23 20:27:04 +0000222Let's try a simple generator:
223
224 >>> def f():
225 ... yield 1
226 ... yield 2
227
Tim Petersb9e9ff12001-06-24 03:44:52 +0000228 >>> for i in f():
Guido van Rossum7131f842007-02-09 20:13:25 +0000229 ... print(i)
Tim Petersb9e9ff12001-06-24 03:44:52 +0000230 1
231 2
Tim Peters1def3512001-06-23 20:27:04 +0000232 >>> g = f()
Georg Brandla18af4e2007-04-21 15:47:16 +0000233 >>> next(g)
Tim Peters1def3512001-06-23 20:27:04 +0000234 1
Georg Brandla18af4e2007-04-21 15:47:16 +0000235 >>> next(g)
Tim Peters1def3512001-06-23 20:27:04 +0000236 2
Tim Petersea2e97a2001-06-24 07:10:02 +0000237
Tim Peters2106ef02001-06-25 01:30:12 +0000238"Falling off the end" stops the generator:
Tim Petersea2e97a2001-06-24 07:10:02 +0000239
Georg Brandla18af4e2007-04-21 15:47:16 +0000240 >>> next(g)
Tim Peters1def3512001-06-23 20:27:04 +0000241 Traceback (most recent call last):
242 File "<stdin>", line 1, in ?
243 File "<stdin>", line 2, in g
244 StopIteration
245
Tim Petersea2e97a2001-06-24 07:10:02 +0000246"return" also stops the generator:
Tim Peters1def3512001-06-23 20:27:04 +0000247
248 >>> def f():
249 ... yield 1
250 ... return
251 ... yield 2 # never reached
252 ...
253 >>> g = f()
Georg Brandla18af4e2007-04-21 15:47:16 +0000254 >>> next(g)
Tim Peters1def3512001-06-23 20:27:04 +0000255 1
Georg Brandla18af4e2007-04-21 15:47:16 +0000256 >>> next(g)
Tim Peters1def3512001-06-23 20:27:04 +0000257 Traceback (most recent call last):
258 File "<stdin>", line 1, in ?
259 File "<stdin>", line 3, in f
260 StopIteration
Georg Brandla18af4e2007-04-21 15:47:16 +0000261 >>> next(g) # once stopped, can't be resumed
Tim Peters1def3512001-06-23 20:27:04 +0000262 Traceback (most recent call last):
263 File "<stdin>", line 1, in ?
264 StopIteration
265
266"raise StopIteration" stops the generator too:
267
268 >>> def f():
269 ... yield 1
Tim Peters34463652001-07-12 22:43:41 +0000270 ... raise StopIteration
Tim Peters1def3512001-06-23 20:27:04 +0000271 ... yield 2 # never reached
272 ...
273 >>> g = f()
Georg Brandla18af4e2007-04-21 15:47:16 +0000274 >>> next(g)
Tim Peters1def3512001-06-23 20:27:04 +0000275 1
Georg Brandla18af4e2007-04-21 15:47:16 +0000276 >>> next(g)
Tim Peters1def3512001-06-23 20:27:04 +0000277 Traceback (most recent call last):
278 File "<stdin>", line 1, in ?
279 StopIteration
Georg Brandla18af4e2007-04-21 15:47:16 +0000280 >>> next(g)
Tim Peters1def3512001-06-23 20:27:04 +0000281 Traceback (most recent call last):
282 File "<stdin>", line 1, in ?
283 StopIteration
284
285However, they are not exactly equivalent:
286
287 >>> def g1():
288 ... try:
289 ... return
290 ... except:
291 ... yield 1
292 ...
293 >>> list(g1())
294 []
295
296 >>> def g2():
297 ... try:
298 ... raise StopIteration
299 ... except:
300 ... yield 42
Guido van Rossum7131f842007-02-09 20:13:25 +0000301 >>> print(list(g2()))
Tim Peters1def3512001-06-23 20:27:04 +0000302 [42]
303
304This may be surprising at first:
305
306 >>> def g3():
307 ... try:
308 ... return
309 ... finally:
310 ... yield 1
311 ...
312 >>> list(g3())
313 [1]
314
315Let's create an alternate range() function implemented as a generator:
316
317 >>> def yrange(n):
318 ... for i in range(n):
319 ... yield i
320 ...
321 >>> list(yrange(5))
322 [0, 1, 2, 3, 4]
323
324Generators always return to the most recent caller:
325
326 >>> def creator():
327 ... r = yrange(5)
Georg Brandla18af4e2007-04-21 15:47:16 +0000328 ... print("creator", next(r))
Tim Peters1def3512001-06-23 20:27:04 +0000329 ... return r
Guido van Rossumc420b2f2007-02-09 22:09:01 +0000330 ...
Tim Peters1def3512001-06-23 20:27:04 +0000331 >>> def caller():
332 ... r = creator()
333 ... for i in r:
Guido van Rossum7131f842007-02-09 20:13:25 +0000334 ... print("caller", i)
Guido van Rossumc420b2f2007-02-09 22:09:01 +0000335 ...
Tim Peters1def3512001-06-23 20:27:04 +0000336 >>> caller()
337 creator 0
338 caller 1
339 caller 2
340 caller 3
341 caller 4
342
343Generators can call other generators:
344
345 >>> def zrange(n):
346 ... for i in yrange(n):
347 ... yield i
348 ...
349 >>> list(zrange(5))
350 [0, 1, 2, 3, 4]
351
352"""
353
Tim Peters6ba5f792001-06-23 20:45:43 +0000354# The examples from PEP 255.
355
356pep_tests = """
357
Tim Peterse5614632001-08-15 04:41:19 +0000358Specification: Yield
359
360 Restriction: A generator cannot be resumed while it is actively
361 running:
362
363 >>> def g():
Georg Brandla18af4e2007-04-21 15:47:16 +0000364 ... i = next(me)
Tim Peterse5614632001-08-15 04:41:19 +0000365 ... yield i
366 >>> me = g()
Georg Brandla18af4e2007-04-21 15:47:16 +0000367 >>> next(me)
Tim Peterse5614632001-08-15 04:41:19 +0000368 Traceback (most recent call last):
369 ...
370 File "<string>", line 2, in g
371 ValueError: generator already executing
372
Tim Peters6ba5f792001-06-23 20:45:43 +0000373Specification: Return
374
375 Note that return isn't always equivalent to raising StopIteration: the
376 difference lies in how enclosing try/except constructs are treated.
377 For example,
378
379 >>> def f1():
380 ... try:
381 ... return
382 ... except:
383 ... yield 1
Guido van Rossum7131f842007-02-09 20:13:25 +0000384 >>> print(list(f1()))
Tim Peters6ba5f792001-06-23 20:45:43 +0000385 []
386
387 because, as in any function, return simply exits, but
388
389 >>> def f2():
390 ... try:
391 ... raise StopIteration
392 ... except:
393 ... yield 42
Guido van Rossum7131f842007-02-09 20:13:25 +0000394 >>> print(list(f2()))
Tim Peters6ba5f792001-06-23 20:45:43 +0000395 [42]
396
397 because StopIteration is captured by a bare "except", as is any
398 exception.
399
400Specification: Generators and Exception Propagation
401
402 >>> def f():
Tim Peters3caca232001-12-06 06:23:26 +0000403 ... return 1//0
Tim Peters6ba5f792001-06-23 20:45:43 +0000404 >>> def g():
405 ... yield f() # the zero division exception propagates
406 ... yield 42 # and we'll never get here
407 >>> k = g()
Georg Brandla18af4e2007-04-21 15:47:16 +0000408 >>> next(k)
Tim Peters6ba5f792001-06-23 20:45:43 +0000409 Traceback (most recent call last):
410 File "<stdin>", line 1, in ?
411 File "<stdin>", line 2, in g
412 File "<stdin>", line 2, in f
413 ZeroDivisionError: integer division or modulo by zero
Georg Brandla18af4e2007-04-21 15:47:16 +0000414 >>> next(k) # and the generator cannot be resumed
Tim Peters6ba5f792001-06-23 20:45:43 +0000415 Traceback (most recent call last):
416 File "<stdin>", line 1, in ?
417 StopIteration
418 >>>
419
420Specification: Try/Except/Finally
421
422 >>> def f():
423 ... try:
424 ... yield 1
425 ... try:
426 ... yield 2
Tim Peters3caca232001-12-06 06:23:26 +0000427 ... 1//0
Tim Peters6ba5f792001-06-23 20:45:43 +0000428 ... yield 3 # never get here
429 ... except ZeroDivisionError:
430 ... yield 4
431 ... yield 5
432 ... raise
433 ... except:
434 ... yield 6
435 ... yield 7 # the "raise" above stops this
436 ... except:
437 ... yield 8
438 ... yield 9
439 ... try:
440 ... x = 12
441 ... finally:
442 ... yield 10
443 ... yield 11
Guido van Rossum7131f842007-02-09 20:13:25 +0000444 >>> print(list(f()))
Tim Peters6ba5f792001-06-23 20:45:43 +0000445 [1, 2, 4, 5, 8, 9, 10, 11]
446 >>>
447
Tim Peters6ba5f792001-06-23 20:45:43 +0000448Guido's binary tree example.
449
450 >>> # A binary tree class.
451 >>> class Tree:
452 ...
453 ... def __init__(self, label, left=None, right=None):
454 ... self.label = label
455 ... self.left = left
456 ... self.right = right
457 ...
458 ... def __repr__(self, level=0, indent=" "):
Walter Dörwald70a6b492004-02-12 17:35:32 +0000459 ... s = level*indent + repr(self.label)
Tim Peters6ba5f792001-06-23 20:45:43 +0000460 ... if self.left:
461 ... s = s + "\\n" + self.left.__repr__(level+1, indent)
462 ... if self.right:
463 ... s = s + "\\n" + self.right.__repr__(level+1, indent)
464 ... return s
465 ...
466 ... def __iter__(self):
467 ... return inorder(self)
468
469 >>> # Create a Tree from a list.
470 >>> def tree(list):
471 ... n = len(list)
472 ... if n == 0:
473 ... return []
Tim Peters3caca232001-12-06 06:23:26 +0000474 ... i = n // 2
Tim Peters6ba5f792001-06-23 20:45:43 +0000475 ... return Tree(list[i], tree(list[:i]), tree(list[i+1:]))
476
477 >>> # Show it off: create a tree.
478 >>> t = tree("ABCDEFGHIJKLMNOPQRSTUVWXYZ")
479
Tim Petersd674e172002-03-10 07:59:13 +0000480 >>> # A recursive generator that generates Tree labels in in-order.
Tim Peters6ba5f792001-06-23 20:45:43 +0000481 >>> def inorder(t):
482 ... if t:
483 ... for x in inorder(t.left):
484 ... yield x
485 ... yield t.label
486 ... for x in inorder(t.right):
487 ... yield x
488
489 >>> # Show it off: create a tree.
Edward Loper103d26e2004-08-09 02:03:30 +0000490 >>> t = tree("ABCDEFGHIJKLMNOPQRSTUVWXYZ")
491 >>> # Print the nodes of the tree in in-order.
492 >>> for x in t:
Guido van Rossumc420b2f2007-02-09 22:09:01 +0000493 ... print(' '+x, end='')
494 A B C D E F G H I J K L M N O P Q R S T U V W X Y Z
Tim Peters6ba5f792001-06-23 20:45:43 +0000495
496 >>> # A non-recursive generator.
497 >>> def inorder(node):
498 ... stack = []
499 ... while node:
500 ... while node.left:
501 ... stack.append(node)
502 ... node = node.left
503 ... yield node.label
504 ... while not node.right:
505 ... try:
506 ... node = stack.pop()
507 ... except IndexError:
508 ... return
509 ... yield node.label
510 ... node = node.right
511
512 >>> # Exercise the non-recursive generator.
513 >>> for x in t:
Guido van Rossumc420b2f2007-02-09 22:09:01 +0000514 ... print(' '+x, end='')
515 A B C D E F G H I J K L M N O P Q R S T U V W X Y Z
Tim Peters6ba5f792001-06-23 20:45:43 +0000516
517"""
518
Tim Petersb2bc6a92001-06-24 10:14:27 +0000519# Examples from Iterator-List and Python-Dev and c.l.py.
Tim Peters6ba5f792001-06-23 20:45:43 +0000520
521email_tests = """
522
523The difference between yielding None and returning it.
524
525>>> def g():
526... for i in range(3):
527... yield None
528... yield None
529... return
530>>> list(g())
531[None, None, None, None]
532
533Ensure that explicitly raising StopIteration acts like any other exception
534in try/except, not like a return.
535
536>>> def g():
537... yield 1
538... try:
539... raise StopIteration
540... except:
541... yield 2
542... yield 3
543>>> list(g())
544[1, 2, 3]
Tim Petersb9e9ff12001-06-24 03:44:52 +0000545
Tim Petersb2bc6a92001-06-24 10:14:27 +0000546Next one was posted to c.l.py.
547
548>>> def gcomb(x, k):
549... "Generate all combinations of k elements from list x."
550...
551... if k > len(x):
552... return
553... if k == 0:
554... yield []
555... else:
556... first, rest = x[0], x[1:]
557... # A combination does or doesn't contain first.
558... # If it does, the remainder is a k-1 comb of rest.
559... for c in gcomb(rest, k-1):
560... c.insert(0, first)
561... yield c
562... # If it doesn't contain first, it's a k comb of rest.
563... for c in gcomb(rest, k):
564... yield c
565
Guido van Rossum805365e2007-05-07 22:24:25 +0000566>>> seq = list(range(1, 5))
Tim Petersb2bc6a92001-06-24 10:14:27 +0000567>>> for k in range(len(seq) + 2):
Guido van Rossum7131f842007-02-09 20:13:25 +0000568... print("%d-combs of %s:" % (k, seq))
Tim Petersb2bc6a92001-06-24 10:14:27 +0000569... for c in gcomb(seq, k):
Guido van Rossum7131f842007-02-09 20:13:25 +0000570... print(" ", c)
Tim Petersb2bc6a92001-06-24 10:14:27 +00005710-combs of [1, 2, 3, 4]:
572 []
5731-combs of [1, 2, 3, 4]:
574 [1]
575 [2]
576 [3]
577 [4]
5782-combs of [1, 2, 3, 4]:
579 [1, 2]
580 [1, 3]
581 [1, 4]
582 [2, 3]
583 [2, 4]
584 [3, 4]
5853-combs of [1, 2, 3, 4]:
586 [1, 2, 3]
587 [1, 2, 4]
588 [1, 3, 4]
589 [2, 3, 4]
5904-combs of [1, 2, 3, 4]:
591 [1, 2, 3, 4]
5925-combs of [1, 2, 3, 4]:
Tim Peters3e7b1a02001-06-25 19:46:25 +0000593
Tim Peterse77f2e22001-06-26 22:24:51 +0000594From the Iterators list, about the types of these things.
Tim Peters3e7b1a02001-06-25 19:46:25 +0000595
596>>> def g():
597... yield 1
598...
599>>> type(g)
Martin v. Löwis250ad612008-04-07 05:43:42 +0000600<class 'function'>
Tim Peters3e7b1a02001-06-25 19:46:25 +0000601>>> i = g()
602>>> type(i)
Martin v. Löwis250ad612008-04-07 05:43:42 +0000603<class 'generator'>
Tim Peters5d2b77c2001-09-03 05:47:38 +0000604>>> [s for s in dir(i) if not s.startswith('_')]
Christian Heimesaf98da12008-01-27 15:18:18 +0000605['close', 'gi_code', 'gi_frame', 'gi_running', 'send', 'throw']
Serhiy Storchaka9a11f172013-01-31 16:11:04 +0200606>>> from test.support import HAVE_DOCSTRINGS
Larry Hastings581ee362014-01-28 05:00:08 -0800607>>> print(i.__next__.__doc__ if HAVE_DOCSTRINGS else 'Implement next(self).')
608Implement next(self).
Tim Peters3e7b1a02001-06-25 19:46:25 +0000609>>> iter(i) is i
Guido van Rossum77f6a652002-04-03 22:41:51 +0000610True
Tim Peters3e7b1a02001-06-25 19:46:25 +0000611>>> import types
612>>> isinstance(i, types.GeneratorType)
Guido van Rossum77f6a652002-04-03 22:41:51 +0000613True
Tim Peterse77f2e22001-06-26 22:24:51 +0000614
615And more, added later.
616
617>>> i.gi_running
6180
619>>> type(i.gi_frame)
Martin v. Löwis250ad612008-04-07 05:43:42 +0000620<class 'frame'>
Tim Peterse77f2e22001-06-26 22:24:51 +0000621>>> i.gi_running = 42
622Traceback (most recent call last):
623 ...
Collin Winter42dae6a2007-03-28 21:44:53 +0000624AttributeError: readonly attribute
Tim Peterse77f2e22001-06-26 22:24:51 +0000625>>> def g():
626... yield me.gi_running
627>>> me = g()
628>>> me.gi_running
6290
Georg Brandla18af4e2007-04-21 15:47:16 +0000630>>> next(me)
Tim Peterse77f2e22001-06-26 22:24:51 +00006311
632>>> me.gi_running
6330
Tim Peters35302662001-07-02 01:38:33 +0000634
635A clever union-find implementation from c.l.py, due to David Eppstein.
636Sent: Friday, June 29, 2001 12:16 PM
637To: python-list@python.org
638Subject: Re: PEP 255: Simple Generators
639
640>>> class disjointSet:
641... def __init__(self, name):
642... self.name = name
643... self.parent = None
644... self.generator = self.generate()
645...
646... def generate(self):
647... while not self.parent:
648... yield self
649... for x in self.parent.generator:
650... yield x
651...
652... def find(self):
Georg Brandla18af4e2007-04-21 15:47:16 +0000653... return next(self.generator)
Tim Peters35302662001-07-02 01:38:33 +0000654...
655... def union(self, parent):
656... if self.parent:
657... raise ValueError("Sorry, I'm not a root!")
658... self.parent = parent
659...
660... def __str__(self):
661... return self.name
Tim Peters35302662001-07-02 01:38:33 +0000662
663>>> names = "ABCDEFGHIJKLM"
664>>> sets = [disjointSet(name) for name in names]
665>>> roots = sets[:]
666
667>>> import random
Raymond Hettinger28de64f2008-01-13 23:40:30 +0000668>>> gen = random.Random(42)
Tim Peters35302662001-07-02 01:38:33 +0000669>>> while 1:
670... for s in sets:
Guido van Rossumc420b2f2007-02-09 22:09:01 +0000671... print(" %s->%s" % (s, s.find()), end='')
Guido van Rossum7131f842007-02-09 20:13:25 +0000672... print()
Tim Peters35302662001-07-02 01:38:33 +0000673... if len(roots) > 1:
Raymond Hettingerdd24a9f2002-12-30 00:46:09 +0000674... s1 = gen.choice(roots)
Tim Peters35302662001-07-02 01:38:33 +0000675... roots.remove(s1)
Raymond Hettingerdd24a9f2002-12-30 00:46:09 +0000676... s2 = gen.choice(roots)
Tim Peters35302662001-07-02 01:38:33 +0000677... s1.union(s2)
Guido van Rossum7131f842007-02-09 20:13:25 +0000678... print("merged", s1, "into", s2)
Tim Peters35302662001-07-02 01:38:33 +0000679... else:
680... break
Guido van Rossumc420b2f2007-02-09 22:09:01 +0000681 A->A B->B C->C D->D E->E F->F G->G H->H I->I J->J K->K L->L M->M
Raymond Hettingerc585eec2010-09-07 15:00:15 +0000682merged K into B
683 A->A B->B C->C D->D E->E F->F G->G H->H I->I J->J K->B L->L M->M
Raymond Hettinger28de64f2008-01-13 23:40:30 +0000684merged A into F
Raymond Hettingerc585eec2010-09-07 15:00:15 +0000685 A->F B->B C->C D->D E->E F->F G->G H->H I->I J->J K->B L->L M->M
686merged E into F
687 A->F B->B C->C D->D E->F F->F G->G H->H I->I J->J K->B L->L M->M
688merged D into C
689 A->F B->B C->C D->C E->F F->F G->G H->H I->I J->J K->B L->L M->M
Raymond Hettinger28de64f2008-01-13 23:40:30 +0000690merged M into C
Raymond Hettingerc585eec2010-09-07 15:00:15 +0000691 A->F B->B C->C D->C E->F F->F G->G H->H I->I J->J K->B L->L M->C
692merged J into B
693 A->F B->B C->C D->C E->F F->F G->G H->H I->I J->B K->B L->L M->C
694merged B into C
695 A->F B->C C->C D->C E->F F->F G->G H->H I->I J->C K->C L->L M->C
696merged F into G
697 A->G B->C C->C D->C E->G F->G G->G H->H I->I J->C K->C L->L M->C
698merged L into C
699 A->G B->C C->C D->C E->G F->G G->G H->H I->I J->C K->C L->C M->C
700merged G into I
701 A->I B->C C->C D->C E->I F->I G->I H->H I->I J->C K->C L->C M->C
702merged I into H
703 A->H B->C C->C D->C E->H F->H G->H H->H I->H J->C K->C L->C M->C
704merged C into H
705 A->H B->H C->H D->H E->H F->H G->H H->H I->H J->H K->H L->H M->H
Phillip J. Eby0d6615f2005-08-02 00:46:46 +0000706
Tim Peters6ba5f792001-06-23 20:45:43 +0000707"""
Barry Warsaw04f357c2002-07-23 19:04:11 +0000708# Emacs turd '
Tim Peters6ba5f792001-06-23 20:45:43 +0000709
Tim Peters0f9da0a2001-06-23 21:01:47 +0000710# Fun tests (for sufficiently warped notions of "fun").
711
712fun_tests = """
713
714Build up to a recursive Sieve of Eratosthenes generator.
715
716>>> def firstn(g, n):
Georg Brandla18af4e2007-04-21 15:47:16 +0000717... return [next(g) for i in range(n)]
Tim Peters0f9da0a2001-06-23 21:01:47 +0000718
719>>> def intsfrom(i):
720... while 1:
721... yield i
722... i += 1
723
724>>> firstn(intsfrom(5), 7)
725[5, 6, 7, 8, 9, 10, 11]
726
727>>> def exclude_multiples(n, ints):
728... for i in ints:
729... if i % n:
730... yield i
731
732>>> firstn(exclude_multiples(3, intsfrom(1)), 6)
733[1, 2, 4, 5, 7, 8]
734
735>>> def sieve(ints):
Georg Brandla18af4e2007-04-21 15:47:16 +0000736... prime = next(ints)
Tim Peters0f9da0a2001-06-23 21:01:47 +0000737... yield prime
738... not_divisible_by_prime = exclude_multiples(prime, ints)
739... for p in sieve(not_divisible_by_prime):
740... yield p
741
742>>> primes = sieve(intsfrom(2))
743>>> firstn(primes, 20)
744[2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71]
Tim Petersb9e9ff12001-06-24 03:44:52 +0000745
Tim Petersf6ed0742001-06-27 07:17:57 +0000746
Tim Petersb9e9ff12001-06-24 03:44:52 +0000747Another famous problem: generate all integers of the form
748 2**i * 3**j * 5**k
749in increasing order, where i,j,k >= 0. Trickier than it may look at first!
750Try writing it without generators, and correctly, and without generating
7513 internal results for each result output.
752
753>>> def times(n, g):
754... for i in g:
755... yield n * i
756>>> firstn(times(10, intsfrom(1)), 10)
757[10, 20, 30, 40, 50, 60, 70, 80, 90, 100]
758
759>>> def merge(g, h):
Georg Brandla18af4e2007-04-21 15:47:16 +0000760... ng = next(g)
761... nh = next(h)
Tim Petersb9e9ff12001-06-24 03:44:52 +0000762... while 1:
763... if ng < nh:
764... yield ng
Georg Brandla18af4e2007-04-21 15:47:16 +0000765... ng = next(g)
Tim Petersb9e9ff12001-06-24 03:44:52 +0000766... elif ng > nh:
767... yield nh
Georg Brandla18af4e2007-04-21 15:47:16 +0000768... nh = next(h)
Tim Petersb9e9ff12001-06-24 03:44:52 +0000769... else:
770... yield ng
Georg Brandla18af4e2007-04-21 15:47:16 +0000771... ng = next(g)
772... nh = next(h)
Tim Petersb9e9ff12001-06-24 03:44:52 +0000773
Tim Petersf6ed0742001-06-27 07:17:57 +0000774The following works, but is doing a whale of a lot of redundant work --
775it's not clear how to get the internal uses of m235 to share a single
776generator. Note that me_times2 (etc) each need to see every element in the
777result sequence. So this is an example where lazy lists are more natural
778(you can look at the head of a lazy list any number of times).
Tim Petersb9e9ff12001-06-24 03:44:52 +0000779
780>>> def m235():
781... yield 1
782... me_times2 = times(2, m235())
783... me_times3 = times(3, m235())
784... me_times5 = times(5, m235())
785... for i in merge(merge(me_times2,
786... me_times3),
787... me_times5):
788... yield i
789
Tim Petersf6ed0742001-06-27 07:17:57 +0000790Don't print "too many" of these -- the implementation above is extremely
791inefficient: each call of m235() leads to 3 recursive calls, and in
792turn each of those 3 more, and so on, and so on, until we've descended
793enough levels to satisfy the print stmts. Very odd: when I printed 5
794lines of results below, this managed to screw up Win98's malloc in "the
795usual" way, i.e. the heap grew over 4Mb so Win98 started fragmenting
796address space, and it *looked* like a very slow leak.
797
Tim Petersb9e9ff12001-06-24 03:44:52 +0000798>>> result = m235()
Tim Petersf6ed0742001-06-27 07:17:57 +0000799>>> for i in range(3):
Guido van Rossum7131f842007-02-09 20:13:25 +0000800... print(firstn(result, 15))
Tim Petersb9e9ff12001-06-24 03:44:52 +0000801[1, 2, 3, 4, 5, 6, 8, 9, 10, 12, 15, 16, 18, 20, 24]
802[25, 27, 30, 32, 36, 40, 45, 48, 50, 54, 60, 64, 72, 75, 80]
803[81, 90, 96, 100, 108, 120, 125, 128, 135, 144, 150, 160, 162, 180, 192]
Tim Petersee309272001-06-24 05:47:06 +0000804
805Heh. Here's one way to get a shared list, complete with an excruciating
806namespace renaming trick. The *pretty* part is that the times() and merge()
807functions can be reused as-is, because they only assume their stream
808arguments are iterable -- a LazyList is the same as a generator to times().
809
810>>> class LazyList:
811... def __init__(self, g):
812... self.sofar = []
Georg Brandla18af4e2007-04-21 15:47:16 +0000813... self.fetch = g.__next__
Tim Petersee309272001-06-24 05:47:06 +0000814...
815... def __getitem__(self, i):
816... sofar, fetch = self.sofar, self.fetch
817... while i >= len(sofar):
818... sofar.append(fetch())
819... return sofar[i]
820
821>>> def m235():
822... yield 1
Tim Petersea2e97a2001-06-24 07:10:02 +0000823... # Gack: m235 below actually refers to a LazyList.
Tim Petersee309272001-06-24 05:47:06 +0000824... me_times2 = times(2, m235)
825... me_times3 = times(3, m235)
826... me_times5 = times(5, m235)
827... for i in merge(merge(me_times2,
828... me_times3),
829... me_times5):
830... yield i
831
Tim Petersf6ed0742001-06-27 07:17:57 +0000832Print as many of these as you like -- *this* implementation is memory-
Neil Schemenauerb20e9db2001-07-12 13:26:41 +0000833efficient.
Tim Petersf6ed0742001-06-27 07:17:57 +0000834
Tim Petersee309272001-06-24 05:47:06 +0000835>>> m235 = LazyList(m235())
836>>> for i in range(5):
Guido van Rossum7131f842007-02-09 20:13:25 +0000837... print([m235[j] for j in range(15*i, 15*(i+1))])
Tim Petersee309272001-06-24 05:47:06 +0000838[1, 2, 3, 4, 5, 6, 8, 9, 10, 12, 15, 16, 18, 20, 24]
839[25, 27, 30, 32, 36, 40, 45, 48, 50, 54, 60, 64, 72, 75, 80]
840[81, 90, 96, 100, 108, 120, 125, 128, 135, 144, 150, 160, 162, 180, 192]
841[200, 216, 225, 240, 243, 250, 256, 270, 288, 300, 320, 324, 360, 375, 384]
842[400, 405, 432, 450, 480, 486, 500, 512, 540, 576, 600, 625, 640, 648, 675]
Tim Petersf6ed0742001-06-27 07:17:57 +0000843
Tim Petersf6ed0742001-06-27 07:17:57 +0000844Ye olde Fibonacci generator, LazyList style.
845
846>>> def fibgen(a, b):
847...
848... def sum(g, h):
849... while 1:
Georg Brandla18af4e2007-04-21 15:47:16 +0000850... yield next(g) + next(h)
Tim Petersf6ed0742001-06-27 07:17:57 +0000851...
852... def tail(g):
Georg Brandla18af4e2007-04-21 15:47:16 +0000853... next(g) # throw first away
Tim Petersf6ed0742001-06-27 07:17:57 +0000854... for x in g:
855... yield x
856...
857... yield a
858... yield b
859... for s in sum(iter(fib),
860... tail(iter(fib))):
861... yield s
862
863>>> fib = LazyList(fibgen(1, 2))
864>>> firstn(iter(fib), 17)
865[1, 2, 3, 5, 8, 13, 21, 34, 55, 89, 144, 233, 377, 610, 987, 1597, 2584]
Georg Brandl52715f62005-08-24 09:02:29 +0000866
867
868Running after your tail with itertools.tee (new in version 2.4)
869
870The algorithms "m235" (Hamming) and Fibonacci presented above are both
871examples of a whole family of FP (functional programming) algorithms
872where a function produces and returns a list while the production algorithm
873suppose the list as already produced by recursively calling itself.
874For these algorithms to work, they must:
875
876- produce at least a first element without presupposing the existence of
877 the rest of the list
878- produce their elements in a lazy manner
879
880To work efficiently, the beginning of the list must not be recomputed over
881and over again. This is ensured in most FP languages as a built-in feature.
882In python, we have to explicitly maintain a list of already computed results
883and abandon genuine recursivity.
884
885This is what had been attempted above with the LazyList class. One problem
886with that class is that it keeps a list of all of the generated results and
887therefore continually grows. This partially defeats the goal of the generator
888concept, viz. produce the results only as needed instead of producing them
889all and thereby wasting memory.
890
891Thanks to itertools.tee, it is now clear "how to get the internal uses of
892m235 to share a single generator".
893
894>>> from itertools import tee
895>>> def m235():
896... def _m235():
897... yield 1
898... for n in merge(times(2, m2),
899... merge(times(3, m3),
900... times(5, m5))):
901... yield n
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000902... m1 = _m235()
903... m2, m3, m5, mRes = tee(m1, 4)
Georg Brandl52715f62005-08-24 09:02:29 +0000904... return mRes
905
906>>> it = m235()
907>>> for i in range(5):
Guido van Rossum7131f842007-02-09 20:13:25 +0000908... print(firstn(it, 15))
Georg Brandl52715f62005-08-24 09:02:29 +0000909[1, 2, 3, 4, 5, 6, 8, 9, 10, 12, 15, 16, 18, 20, 24]
910[25, 27, 30, 32, 36, 40, 45, 48, 50, 54, 60, 64, 72, 75, 80]
911[81, 90, 96, 100, 108, 120, 125, 128, 135, 144, 150, 160, 162, 180, 192]
912[200, 216, 225, 240, 243, 250, 256, 270, 288, 300, 320, 324, 360, 375, 384]
913[400, 405, 432, 450, 480, 486, 500, 512, 540, 576, 600, 625, 640, 648, 675]
914
915The "tee" function does just what we want. It internally keeps a generated
916result for as long as it has not been "consumed" from all of the duplicated
917iterators, whereupon it is deleted. You can therefore print the hamming
918sequence during hours without increasing memory usage, or very little.
919
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000920The beauty of it is that recursive running-after-their-tail FP algorithms
Georg Brandl52715f62005-08-24 09:02:29 +0000921are quite straightforwardly expressed with this Python idiom.
922
Georg Brandl52715f62005-08-24 09:02:29 +0000923Ye olde Fibonacci generator, tee style.
924
925>>> def fib():
Tim Peters9e34c042005-08-26 15:20:46 +0000926...
Georg Brandl52715f62005-08-24 09:02:29 +0000927... def _isum(g, h):
928... while 1:
Georg Brandla18af4e2007-04-21 15:47:16 +0000929... yield next(g) + next(h)
Tim Peters9e34c042005-08-26 15:20:46 +0000930...
Georg Brandl52715f62005-08-24 09:02:29 +0000931... def _fib():
932... yield 1
933... yield 2
Georg Brandla18af4e2007-04-21 15:47:16 +0000934... next(fibTail) # throw first away
Georg Brandl52715f62005-08-24 09:02:29 +0000935... for res in _isum(fibHead, fibTail):
936... yield res
Tim Peters9e34c042005-08-26 15:20:46 +0000937...
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000938... realfib = _fib()
939... fibHead, fibTail, fibRes = tee(realfib, 3)
Georg Brandl52715f62005-08-24 09:02:29 +0000940... return fibRes
941
942>>> firstn(fib(), 17)
943[1, 2, 3, 5, 8, 13, 21, 34, 55, 89, 144, 233, 377, 610, 987, 1597, 2584]
944
Tim Peters0f9da0a2001-06-23 21:01:47 +0000945"""
946
Tim Petersb6c3cea2001-06-26 03:36:28 +0000947# syntax_tests mostly provokes SyntaxErrors. Also fiddling with #if 0
948# hackery.
Tim Petersee309272001-06-24 05:47:06 +0000949
Tim Petersea2e97a2001-06-24 07:10:02 +0000950syntax_tests = """
951
Phillip J. Eby0d6615f2005-08-02 00:46:46 +0000952These are fine:
Tim Petersea2e97a2001-06-24 07:10:02 +0000953
954>>> def f():
955... yield 1
956... return
957
Tim Petersaef8cfa2004-08-27 15:12:49 +0000958>>> def f():
Tim Petersea2e97a2001-06-24 07:10:02 +0000959... try:
960... yield 1
961... finally:
962... pass
Tim Petersea2e97a2001-06-24 07:10:02 +0000963
Tim Petersaef8cfa2004-08-27 15:12:49 +0000964>>> def f():
Tim Petersea2e97a2001-06-24 07:10:02 +0000965... try:
966... try:
Tim Peters3caca232001-12-06 06:23:26 +0000967... 1//0
Tim Petersea2e97a2001-06-24 07:10:02 +0000968... except ZeroDivisionError:
Tim Peters536cf992005-12-25 23:18:31 +0000969... yield 666
Tim Petersea2e97a2001-06-24 07:10:02 +0000970... except:
971... pass
972... finally:
973... pass
Tim Petersea2e97a2001-06-24 07:10:02 +0000974
975>>> def f():
976... try:
977... try:
978... yield 12
Tim Peters3caca232001-12-06 06:23:26 +0000979... 1//0
Tim Petersea2e97a2001-06-24 07:10:02 +0000980... except ZeroDivisionError:
981... yield 666
982... except:
983... try:
984... x = 12
985... finally:
986... yield 12
987... except:
988... return
989>>> list(f())
990[12, 666]
Tim Petersb6c3cea2001-06-26 03:36:28 +0000991
992>>> def f():
Tim Peters08a898f2001-06-28 01:52:22 +0000993... yield
Phillip J. Eby0d6615f2005-08-02 00:46:46 +0000994>>> type(f())
Martin v. Löwis250ad612008-04-07 05:43:42 +0000995<class 'generator'>
Phillip J. Eby0d6615f2005-08-02 00:46:46 +0000996
Tim Peters08a898f2001-06-28 01:52:22 +0000997
998>>> def f():
999... if 0:
1000... yield
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001001>>> type(f())
Martin v. Löwis250ad612008-04-07 05:43:42 +00001002<class 'generator'>
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001003
Tim Peters08a898f2001-06-28 01:52:22 +00001004
1005>>> def f():
Tim Petersb6c3cea2001-06-26 03:36:28 +00001006... if 0:
1007... yield 1
1008>>> type(f())
Martin v. Löwis250ad612008-04-07 05:43:42 +00001009<class 'generator'>
Tim Petersb6c3cea2001-06-26 03:36:28 +00001010
1011>>> def f():
1012... if "":
1013... yield None
1014>>> type(f())
Martin v. Löwis250ad612008-04-07 05:43:42 +00001015<class 'generator'>
Tim Petersb6c3cea2001-06-26 03:36:28 +00001016
1017>>> def f():
1018... return
1019... try:
1020... if x==4:
1021... pass
1022... elif 0:
1023... try:
Tim Peters3caca232001-12-06 06:23:26 +00001024... 1//0
Tim Petersb6c3cea2001-06-26 03:36:28 +00001025... except SyntaxError:
1026... pass
1027... else:
1028... if 0:
1029... while 12:
1030... x += 1
1031... yield 2 # don't blink
1032... f(a, b, c, d, e)
1033... else:
1034... pass
1035... except:
1036... x = 1
1037... return
1038>>> type(f())
Martin v. Löwis250ad612008-04-07 05:43:42 +00001039<class 'generator'>
Tim Petersb6c3cea2001-06-26 03:36:28 +00001040
1041>>> def f():
1042... if 0:
1043... def g():
1044... yield 1
1045...
1046>>> type(f())
Martin v. Löwis250ad612008-04-07 05:43:42 +00001047<class 'NoneType'>
Tim Petersb6c3cea2001-06-26 03:36:28 +00001048
1049>>> def f():
1050... if 0:
1051... class C:
1052... def __init__(self):
1053... yield 1
1054... def f(self):
1055... yield 2
1056>>> type(f())
Martin v. Löwis250ad612008-04-07 05:43:42 +00001057<class 'NoneType'>
Tim Peters08a898f2001-06-28 01:52:22 +00001058
1059>>> def f():
1060... if 0:
1061... return
1062... if 0:
1063... yield 2
1064>>> type(f())
Martin v. Löwis250ad612008-04-07 05:43:42 +00001065<class 'generator'>
Tim Peters08a898f2001-06-28 01:52:22 +00001066
Guido van Rossumc5fe5eb2002-06-12 03:45:21 +00001067This one caused a crash (see SF bug 567538):
1068
1069>>> def f():
1070... for i in range(3):
1071... try:
1072... continue
1073... finally:
1074... yield i
Tim Petersc411dba2002-07-16 21:35:23 +00001075...
Guido van Rossumc5fe5eb2002-06-12 03:45:21 +00001076>>> g = f()
Georg Brandla18af4e2007-04-21 15:47:16 +00001077>>> print(next(g))
Guido van Rossumc5fe5eb2002-06-12 03:45:21 +000010780
Georg Brandla18af4e2007-04-21 15:47:16 +00001079>>> print(next(g))
Guido van Rossumc5fe5eb2002-06-12 03:45:21 +000010801
Georg Brandla18af4e2007-04-21 15:47:16 +00001081>>> print(next(g))
Guido van Rossumc5fe5eb2002-06-12 03:45:21 +000010822
Georg Brandla18af4e2007-04-21 15:47:16 +00001083>>> print(next(g))
Guido van Rossumc5fe5eb2002-06-12 03:45:21 +00001084Traceback (most recent call last):
1085StopIteration
Christian Heimesaf98da12008-01-27 15:18:18 +00001086
1087
1088Test the gi_code attribute
1089
1090>>> def f():
1091... yield 5
1092...
1093>>> g = f()
1094>>> g.gi_code is f.__code__
1095True
1096>>> next(g)
10975
1098>>> next(g)
1099Traceback (most recent call last):
1100StopIteration
1101>>> g.gi_code is f.__code__
1102True
1103
Alexandre Vassalottie9f305f2008-05-16 04:39:54 +00001104
1105Test the __name__ attribute and the repr()
1106
1107>>> def f():
1108... yield 5
1109...
1110>>> g = f()
1111>>> g.__name__
1112'f'
1113>>> repr(g) # doctest: +ELLIPSIS
Alexandre Vassalottibee32532008-05-16 18:15:12 +00001114'<generator object f at ...>'
Benjamin Peterson371ccfb2008-12-27 19:03:36 +00001115
1116Lambdas shouldn't have their usual return behavior.
1117
1118>>> x = lambda: (yield 1)
1119>>> list(x())
1120[1]
1121
1122>>> x = lambda: ((yield 1), (yield 2))
1123>>> list(x())
1124[1, 2]
Tim Petersea2e97a2001-06-24 07:10:02 +00001125"""
1126
Tim Petersbe4f0a72001-06-29 02:41:16 +00001127# conjoin is a simple backtracking generator, named in honor of Icon's
1128# "conjunction" control structure. Pass a list of no-argument functions
1129# that return iterable objects. Easiest to explain by example: assume the
1130# function list [x, y, z] is passed. Then conjoin acts like:
1131#
1132# def g():
1133# values = [None] * 3
1134# for values[0] in x():
1135# for values[1] in y():
1136# for values[2] in z():
1137# yield values
1138#
1139# So some 3-lists of values *may* be generated, each time we successfully
1140# get into the innermost loop. If an iterator fails (is exhausted) before
1141# then, it "backtracks" to get the next value from the nearest enclosing
1142# iterator (the one "to the left"), and starts all over again at the next
1143# slot (pumps a fresh iterator). Of course this is most useful when the
1144# iterators have side-effects, so that which values *can* be generated at
1145# each slot depend on the values iterated at previous slots.
1146
Benjamin Peterson78565b22009-06-28 19:19:51 +00001147def simple_conjoin(gs):
Tim Petersbe4f0a72001-06-29 02:41:16 +00001148
1149 values = [None] * len(gs)
1150
Benjamin Peterson78565b22009-06-28 19:19:51 +00001151 def gen(i):
Tim Petersbe4f0a72001-06-29 02:41:16 +00001152 if i >= len(gs):
1153 yield values
1154 else:
1155 for values[i] in gs[i]():
1156 for x in gen(i+1):
1157 yield x
1158
1159 for x in gen(0):
1160 yield x
1161
Tim Petersc468fd22001-06-30 07:29:44 +00001162# That works fine, but recursing a level and checking i against len(gs) for
1163# each item produced is inefficient. By doing manual loop unrolling across
1164# generator boundaries, it's possible to eliminate most of that overhead.
1165# This isn't worth the bother *in general* for generators, but conjoin() is
1166# a core building block for some CPU-intensive generator applications.
1167
1168def conjoin(gs):
1169
1170 n = len(gs)
1171 values = [None] * n
1172
1173 # Do one loop nest at time recursively, until the # of loop nests
1174 # remaining is divisible by 3.
1175
Benjamin Peterson78565b22009-06-28 19:19:51 +00001176 def gen(i):
Tim Petersc468fd22001-06-30 07:29:44 +00001177 if i >= n:
1178 yield values
1179
1180 elif (n-i) % 3:
1181 ip1 = i+1
1182 for values[i] in gs[i]():
1183 for x in gen(ip1):
1184 yield x
1185
1186 else:
1187 for x in _gen3(i):
1188 yield x
1189
1190 # Do three loop nests at a time, recursing only if at least three more
1191 # remain. Don't call directly: this is an internal optimization for
1192 # gen's use.
1193
Benjamin Peterson78565b22009-06-28 19:19:51 +00001194 def _gen3(i):
Tim Petersc468fd22001-06-30 07:29:44 +00001195 assert i < n and (n-i) % 3 == 0
1196 ip1, ip2, ip3 = i+1, i+2, i+3
1197 g, g1, g2 = gs[i : ip3]
1198
1199 if ip3 >= n:
1200 # These are the last three, so we can yield values directly.
1201 for values[i] in g():
1202 for values[ip1] in g1():
1203 for values[ip2] in g2():
1204 yield values
1205
1206 else:
1207 # At least 6 loop nests remain; peel off 3 and recurse for the
1208 # rest.
1209 for values[i] in g():
1210 for values[ip1] in g1():
1211 for values[ip2] in g2():
1212 for x in _gen3(ip3):
1213 yield x
1214
1215 for x in gen(0):
1216 yield x
1217
unknown31569562001-07-04 22:11:22 +00001218# And one more approach: For backtracking apps like the Knight's Tour
1219# solver below, the number of backtracking levels can be enormous (one
1220# level per square, for the Knight's Tour, so that e.g. a 100x100 board
1221# needs 10,000 levels). In such cases Python is likely to run out of
1222# stack space due to recursion. So here's a recursion-free version of
1223# conjoin too.
1224# NOTE WELL: This allows large problems to be solved with only trivial
1225# demands on stack space. Without explicitly resumable generators, this is
Tim Peters9a8c8e22001-07-13 09:12:12 +00001226# much harder to achieve. OTOH, this is much slower (up to a factor of 2)
1227# than the fancy unrolled recursive conjoin.
unknown31569562001-07-04 22:11:22 +00001228
1229def flat_conjoin(gs): # rename to conjoin to run tests with this instead
1230 n = len(gs)
1231 values = [None] * n
1232 iters = [None] * n
Tim Peters9a8c8e22001-07-13 09:12:12 +00001233 _StopIteration = StopIteration # make local because caught a *lot*
unknown31569562001-07-04 22:11:22 +00001234 i = 0
Tim Peters9a8c8e22001-07-13 09:12:12 +00001235 while 1:
1236 # Descend.
1237 try:
1238 while i < n:
Georg Brandla18af4e2007-04-21 15:47:16 +00001239 it = iters[i] = gs[i]().__next__
Tim Peters9a8c8e22001-07-13 09:12:12 +00001240 values[i] = it()
1241 i += 1
1242 except _StopIteration:
1243 pass
unknown31569562001-07-04 22:11:22 +00001244 else:
Tim Peters9a8c8e22001-07-13 09:12:12 +00001245 assert i == n
1246 yield values
unknown31569562001-07-04 22:11:22 +00001247
Tim Peters9a8c8e22001-07-13 09:12:12 +00001248 # Backtrack until an older iterator can be resumed.
1249 i -= 1
unknown31569562001-07-04 22:11:22 +00001250 while i >= 0:
1251 try:
1252 values[i] = iters[i]()
Tim Peters9a8c8e22001-07-13 09:12:12 +00001253 # Success! Start fresh at next level.
unknown31569562001-07-04 22:11:22 +00001254 i += 1
1255 break
Tim Peters9a8c8e22001-07-13 09:12:12 +00001256 except _StopIteration:
1257 # Continue backtracking.
1258 i -= 1
1259 else:
1260 assert i < 0
1261 break
unknown31569562001-07-04 22:11:22 +00001262
Tim Petersbe4f0a72001-06-29 02:41:16 +00001263# A conjoin-based N-Queens solver.
1264
1265class Queens:
1266 def __init__(self, n):
1267 self.n = n
1268 rangen = range(n)
1269
1270 # Assign a unique int to each column and diagonal.
1271 # columns: n of those, range(n).
1272 # NW-SE diagonals: 2n-1 of these, i-j unique and invariant along
1273 # each, smallest i-j is 0-(n-1) = 1-n, so add n-1 to shift to 0-
1274 # based.
1275 # NE-SW diagonals: 2n-1 of these, i+j unique and invariant along
1276 # each, smallest i+j is 0, largest is 2n-2.
1277
1278 # For each square, compute a bit vector of the columns and
1279 # diagonals it covers, and for each row compute a function that
1280 # generates the possiblities for the columns in that row.
1281 self.rowgenerators = []
1282 for i in rangen:
Guido van Rossume2a383d2007-01-15 16:59:06 +00001283 rowuses = [(1 << j) | # column ordinal
1284 (1 << (n + i-j + n-1)) | # NW-SE ordinal
1285 (1 << (n + 2*n-1 + i+j)) # NE-SW ordinal
Tim Petersbe4f0a72001-06-29 02:41:16 +00001286 for j in rangen]
1287
1288 def rowgen(rowuses=rowuses):
1289 for j in rangen:
1290 uses = rowuses[j]
Tim Petersc468fd22001-06-30 07:29:44 +00001291 if uses & self.used == 0:
1292 self.used |= uses
1293 yield j
1294 self.used &= ~uses
Tim Petersbe4f0a72001-06-29 02:41:16 +00001295
1296 self.rowgenerators.append(rowgen)
1297
1298 # Generate solutions.
1299 def solve(self):
1300 self.used = 0
1301 for row2col in conjoin(self.rowgenerators):
1302 yield row2col
1303
1304 def printsolution(self, row2col):
1305 n = self.n
1306 assert n == len(row2col)
1307 sep = "+" + "-+" * n
Guido van Rossumbe19ed72007-02-09 05:37:30 +00001308 print(sep)
Tim Petersbe4f0a72001-06-29 02:41:16 +00001309 for i in range(n):
1310 squares = [" " for j in range(n)]
1311 squares[row2col[i]] = "Q"
Guido van Rossumbe19ed72007-02-09 05:37:30 +00001312 print("|" + "|".join(squares) + "|")
1313 print(sep)
Tim Petersbe4f0a72001-06-29 02:41:16 +00001314
unknown31569562001-07-04 22:11:22 +00001315# A conjoin-based Knight's Tour solver. This is pretty sophisticated
1316# (e.g., when used with flat_conjoin above, and passing hard=1 to the
1317# constructor, a 200x200 Knight's Tour was found quickly -- note that we're
Tim Peters9a8c8e22001-07-13 09:12:12 +00001318# creating 10s of thousands of generators then!), and is lengthy.
unknown31569562001-07-04 22:11:22 +00001319
1320class Knights:
Tim Peters9a8c8e22001-07-13 09:12:12 +00001321 def __init__(self, m, n, hard=0):
1322 self.m, self.n = m, n
unknown31569562001-07-04 22:11:22 +00001323
Tim Peters9a8c8e22001-07-13 09:12:12 +00001324 # solve() will set up succs[i] to be a list of square #i's
1325 # successors.
1326 succs = self.succs = []
unknown31569562001-07-04 22:11:22 +00001327
Tim Peters9a8c8e22001-07-13 09:12:12 +00001328 # Remove i0 from each of its successor's successor lists, i.e.
1329 # successors can't go back to i0 again. Return 0 if we can
1330 # detect this makes a solution impossible, else return 1.
unknown31569562001-07-04 22:11:22 +00001331
Tim Peters9a8c8e22001-07-13 09:12:12 +00001332 def remove_from_successors(i0, len=len):
unknown31569562001-07-04 22:11:22 +00001333 # If we remove all exits from a free square, we're dead:
1334 # even if we move to it next, we can't leave it again.
1335 # If we create a square with one exit, we must visit it next;
1336 # else somebody else will have to visit it, and since there's
1337 # only one adjacent, there won't be a way to leave it again.
1338 # Finelly, if we create more than one free square with a
1339 # single exit, we can only move to one of them next, leaving
1340 # the other one a dead end.
1341 ne0 = ne1 = 0
1342 for i in succs[i0]:
Tim Peters9a8c8e22001-07-13 09:12:12 +00001343 s = succs[i]
1344 s.remove(i0)
1345 e = len(s)
1346 if e == 0:
1347 ne0 += 1
1348 elif e == 1:
1349 ne1 += 1
unknown31569562001-07-04 22:11:22 +00001350 return ne0 == 0 and ne1 < 2
1351
Tim Peters9a8c8e22001-07-13 09:12:12 +00001352 # Put i0 back in each of its successor's successor lists.
1353
1354 def add_to_successors(i0):
unknown31569562001-07-04 22:11:22 +00001355 for i in succs[i0]:
Tim Peters9a8c8e22001-07-13 09:12:12 +00001356 succs[i].append(i0)
unknown31569562001-07-04 22:11:22 +00001357
1358 # Generate the first move.
1359 def first():
Tim Peters9a8c8e22001-07-13 09:12:12 +00001360 if m < 1 or n < 1:
unknown31569562001-07-04 22:11:22 +00001361 return
1362
unknown31569562001-07-04 22:11:22 +00001363 # Since we're looking for a cycle, it doesn't matter where we
1364 # start. Starting in a corner makes the 2nd move easy.
Tim Peters9a8c8e22001-07-13 09:12:12 +00001365 corner = self.coords2index(0, 0)
1366 remove_from_successors(corner)
unknown31569562001-07-04 22:11:22 +00001367 self.lastij = corner
1368 yield corner
Tim Peters9a8c8e22001-07-13 09:12:12 +00001369 add_to_successors(corner)
unknown31569562001-07-04 22:11:22 +00001370
1371 # Generate the second moves.
1372 def second():
Tim Peters9a8c8e22001-07-13 09:12:12 +00001373 corner = self.coords2index(0, 0)
unknown31569562001-07-04 22:11:22 +00001374 assert self.lastij == corner # i.e., we started in the corner
Tim Peters9a8c8e22001-07-13 09:12:12 +00001375 if m < 3 or n < 3:
unknown31569562001-07-04 22:11:22 +00001376 return
Tim Peters9a8c8e22001-07-13 09:12:12 +00001377 assert len(succs[corner]) == 2
1378 assert self.coords2index(1, 2) in succs[corner]
1379 assert self.coords2index(2, 1) in succs[corner]
unknown31569562001-07-04 22:11:22 +00001380 # Only two choices. Whichever we pick, the other must be the
Tim Peters9a8c8e22001-07-13 09:12:12 +00001381 # square picked on move m*n, as it's the only way to get back
unknown31569562001-07-04 22:11:22 +00001382 # to (0, 0). Save its index in self.final so that moves before
1383 # the last know it must be kept free.
1384 for i, j in (1, 2), (2, 1):
Tim Peters9a8c8e22001-07-13 09:12:12 +00001385 this = self.coords2index(i, j)
1386 final = self.coords2index(3-i, 3-j)
unknown31569562001-07-04 22:11:22 +00001387 self.final = final
unknown31569562001-07-04 22:11:22 +00001388
Tim Peters9a8c8e22001-07-13 09:12:12 +00001389 remove_from_successors(this)
1390 succs[final].append(corner)
unknown31569562001-07-04 22:11:22 +00001391 self.lastij = this
1392 yield this
Tim Peters9a8c8e22001-07-13 09:12:12 +00001393 succs[final].remove(corner)
1394 add_to_successors(this)
unknown31569562001-07-04 22:11:22 +00001395
Tim Peters9a8c8e22001-07-13 09:12:12 +00001396 # Generate moves 3 thru m*n-1.
1397 def advance(len=len):
unknown31569562001-07-04 22:11:22 +00001398 # If some successor has only one exit, must take it.
1399 # Else favor successors with fewer exits.
1400 candidates = []
1401 for i in succs[self.lastij]:
Tim Peters9a8c8e22001-07-13 09:12:12 +00001402 e = len(succs[i])
1403 assert e > 0, "else remove_from_successors() pruning flawed"
1404 if e == 1:
1405 candidates = [(e, i)]
1406 break
1407 candidates.append((e, i))
unknown31569562001-07-04 22:11:22 +00001408 else:
1409 candidates.sort()
1410
1411 for e, i in candidates:
1412 if i != self.final:
Tim Peters9a8c8e22001-07-13 09:12:12 +00001413 if remove_from_successors(i):
unknown31569562001-07-04 22:11:22 +00001414 self.lastij = i
1415 yield i
Tim Peters9a8c8e22001-07-13 09:12:12 +00001416 add_to_successors(i)
unknown31569562001-07-04 22:11:22 +00001417
Tim Peters9a8c8e22001-07-13 09:12:12 +00001418 # Generate moves 3 thru m*n-1. Alternative version using a
unknown31569562001-07-04 22:11:22 +00001419 # stronger (but more expensive) heuristic to order successors.
Tim Peters9a8c8e22001-07-13 09:12:12 +00001420 # Since the # of backtracking levels is m*n, a poor move early on
1421 # can take eons to undo. Smallest square board for which this
1422 # matters a lot is 52x52.
1423 def advance_hard(vmid=(m-1)/2.0, hmid=(n-1)/2.0, len=len):
unknown31569562001-07-04 22:11:22 +00001424 # If some successor has only one exit, must take it.
1425 # Else favor successors with fewer exits.
1426 # Break ties via max distance from board centerpoint (favor
1427 # corners and edges whenever possible).
1428 candidates = []
1429 for i in succs[self.lastij]:
Tim Peters9a8c8e22001-07-13 09:12:12 +00001430 e = len(succs[i])
1431 assert e > 0, "else remove_from_successors() pruning flawed"
1432 if e == 1:
1433 candidates = [(e, 0, i)]
1434 break
1435 i1, j1 = self.index2coords(i)
1436 d = (i1 - vmid)**2 + (j1 - hmid)**2
1437 candidates.append((e, -d, i))
unknown31569562001-07-04 22:11:22 +00001438 else:
1439 candidates.sort()
1440
1441 for e, d, i in candidates:
1442 if i != self.final:
Tim Peters9a8c8e22001-07-13 09:12:12 +00001443 if remove_from_successors(i):
unknown31569562001-07-04 22:11:22 +00001444 self.lastij = i
1445 yield i
Tim Peters9a8c8e22001-07-13 09:12:12 +00001446 add_to_successors(i)
unknown31569562001-07-04 22:11:22 +00001447
1448 # Generate the last move.
1449 def last():
1450 assert self.final in succs[self.lastij]
1451 yield self.final
1452
Tim Peters9a8c8e22001-07-13 09:12:12 +00001453 if m*n < 4:
1454 self.squaregenerators = [first]
unknown31569562001-07-04 22:11:22 +00001455 else:
Tim Peters9a8c8e22001-07-13 09:12:12 +00001456 self.squaregenerators = [first, second] + \
1457 [hard and advance_hard or advance] * (m*n - 3) + \
unknown31569562001-07-04 22:11:22 +00001458 [last]
1459
Tim Peters9a8c8e22001-07-13 09:12:12 +00001460 def coords2index(self, i, j):
1461 assert 0 <= i < self.m
1462 assert 0 <= j < self.n
1463 return i * self.n + j
1464
1465 def index2coords(self, index):
1466 assert 0 <= index < self.m * self.n
1467 return divmod(index, self.n)
1468
1469 def _init_board(self):
1470 succs = self.succs
1471 del succs[:]
1472 m, n = self.m, self.n
1473 c2i = self.coords2index
1474
1475 offsets = [( 1, 2), ( 2, 1), ( 2, -1), ( 1, -2),
1476 (-1, -2), (-2, -1), (-2, 1), (-1, 2)]
1477 rangen = range(n)
1478 for i in range(m):
1479 for j in rangen:
1480 s = [c2i(i+io, j+jo) for io, jo in offsets
1481 if 0 <= i+io < m and
1482 0 <= j+jo < n]
1483 succs.append(s)
1484
unknown31569562001-07-04 22:11:22 +00001485 # Generate solutions.
1486 def solve(self):
Tim Peters9a8c8e22001-07-13 09:12:12 +00001487 self._init_board()
1488 for x in conjoin(self.squaregenerators):
unknown31569562001-07-04 22:11:22 +00001489 yield x
1490
1491 def printsolution(self, x):
Tim Peters9a8c8e22001-07-13 09:12:12 +00001492 m, n = self.m, self.n
1493 assert len(x) == m*n
1494 w = len(str(m*n))
unknown31569562001-07-04 22:11:22 +00001495 format = "%" + str(w) + "d"
1496
Tim Peters9a8c8e22001-07-13 09:12:12 +00001497 squares = [[None] * n for i in range(m)]
unknown31569562001-07-04 22:11:22 +00001498 k = 1
1499 for i in x:
Tim Peters9a8c8e22001-07-13 09:12:12 +00001500 i1, j1 = self.index2coords(i)
unknown31569562001-07-04 22:11:22 +00001501 squares[i1][j1] = format % k
1502 k += 1
1503
1504 sep = "+" + ("-" * w + "+") * n
Guido van Rossumbe19ed72007-02-09 05:37:30 +00001505 print(sep)
Tim Peters9a8c8e22001-07-13 09:12:12 +00001506 for i in range(m):
unknown31569562001-07-04 22:11:22 +00001507 row = squares[i]
Guido van Rossumbe19ed72007-02-09 05:37:30 +00001508 print("|" + "|".join(row) + "|")
1509 print(sep)
unknown31569562001-07-04 22:11:22 +00001510
Tim Petersbe4f0a72001-06-29 02:41:16 +00001511conjoin_tests = """
1512
1513Generate the 3-bit binary numbers in order. This illustrates dumbest-
1514possible use of conjoin, just to generate the full cross-product.
1515
unknown31569562001-07-04 22:11:22 +00001516>>> for c in conjoin([lambda: iter((0, 1))] * 3):
Guido van Rossum7131f842007-02-09 20:13:25 +00001517... print(c)
Tim Petersbe4f0a72001-06-29 02:41:16 +00001518[0, 0, 0]
1519[0, 0, 1]
1520[0, 1, 0]
1521[0, 1, 1]
1522[1, 0, 0]
1523[1, 0, 1]
1524[1, 1, 0]
1525[1, 1, 1]
1526
Tim Petersc468fd22001-06-30 07:29:44 +00001527For efficiency in typical backtracking apps, conjoin() yields the same list
1528object each time. So if you want to save away a full account of its
1529generated sequence, you need to copy its results.
1530
1531>>> def gencopy(iterator):
1532... for x in iterator:
1533... yield x[:]
1534
1535>>> for n in range(10):
unknown31569562001-07-04 22:11:22 +00001536... all = list(gencopy(conjoin([lambda: iter((0, 1))] * n)))
Guido van Rossum7131f842007-02-09 20:13:25 +00001537... print(n, len(all), all[0] == [0] * n, all[-1] == [1] * n)
Guido van Rossum77f6a652002-04-03 22:41:51 +000015380 1 True True
15391 2 True True
15402 4 True True
15413 8 True True
15424 16 True True
15435 32 True True
15446 64 True True
15457 128 True True
15468 256 True True
15479 512 True True
Tim Petersc468fd22001-06-30 07:29:44 +00001548
Tim Petersbe4f0a72001-06-29 02:41:16 +00001549And run an 8-queens solver.
1550
1551>>> q = Queens(8)
1552>>> LIMIT = 2
1553>>> count = 0
1554>>> for row2col in q.solve():
1555... count += 1
1556... if count <= LIMIT:
Guido van Rossum7131f842007-02-09 20:13:25 +00001557... print("Solution", count)
Tim Petersbe4f0a72001-06-29 02:41:16 +00001558... q.printsolution(row2col)
1559Solution 1
1560+-+-+-+-+-+-+-+-+
1561|Q| | | | | | | |
1562+-+-+-+-+-+-+-+-+
1563| | | | |Q| | | |
1564+-+-+-+-+-+-+-+-+
1565| | | | | | | |Q|
1566+-+-+-+-+-+-+-+-+
1567| | | | | |Q| | |
1568+-+-+-+-+-+-+-+-+
1569| | |Q| | | | | |
1570+-+-+-+-+-+-+-+-+
1571| | | | | | |Q| |
1572+-+-+-+-+-+-+-+-+
1573| |Q| | | | | | |
1574+-+-+-+-+-+-+-+-+
1575| | | |Q| | | | |
1576+-+-+-+-+-+-+-+-+
1577Solution 2
1578+-+-+-+-+-+-+-+-+
1579|Q| | | | | | | |
1580+-+-+-+-+-+-+-+-+
1581| | | | | |Q| | |
1582+-+-+-+-+-+-+-+-+
1583| | | | | | | |Q|
1584+-+-+-+-+-+-+-+-+
1585| | |Q| | | | | |
1586+-+-+-+-+-+-+-+-+
1587| | | | | | |Q| |
1588+-+-+-+-+-+-+-+-+
1589| | | |Q| | | | |
1590+-+-+-+-+-+-+-+-+
1591| |Q| | | | | | |
1592+-+-+-+-+-+-+-+-+
1593| | | | |Q| | | |
1594+-+-+-+-+-+-+-+-+
1595
Guido van Rossum7131f842007-02-09 20:13:25 +00001596>>> print(count, "solutions in all.")
Tim Petersbe4f0a72001-06-29 02:41:16 +0000159792 solutions in all.
unknown31569562001-07-04 22:11:22 +00001598
1599And run a Knight's Tour on a 10x10 board. Note that there are about
160020,000 solutions even on a 6x6 board, so don't dare run this to exhaustion.
1601
Tim Peters9a8c8e22001-07-13 09:12:12 +00001602>>> k = Knights(10, 10)
unknown31569562001-07-04 22:11:22 +00001603>>> LIMIT = 2
1604>>> count = 0
1605>>> for x in k.solve():
1606... count += 1
1607... if count <= LIMIT:
Guido van Rossum7131f842007-02-09 20:13:25 +00001608... print("Solution", count)
unknown31569562001-07-04 22:11:22 +00001609... k.printsolution(x)
1610... else:
1611... break
1612Solution 1
1613+---+---+---+---+---+---+---+---+---+---+
1614| 1| 58| 27| 34| 3| 40| 29| 10| 5| 8|
1615+---+---+---+---+---+---+---+---+---+---+
1616| 26| 35| 2| 57| 28| 33| 4| 7| 30| 11|
1617+---+---+---+---+---+---+---+---+---+---+
1618| 59|100| 73| 36| 41| 56| 39| 32| 9| 6|
1619+---+---+---+---+---+---+---+---+---+---+
1620| 74| 25| 60| 55| 72| 37| 42| 49| 12| 31|
1621+---+---+---+---+---+---+---+---+---+---+
1622| 61| 86| 99| 76| 63| 52| 47| 38| 43| 50|
1623+---+---+---+---+---+---+---+---+---+---+
1624| 24| 75| 62| 85| 54| 71| 64| 51| 48| 13|
1625+---+---+---+---+---+---+---+---+---+---+
1626| 87| 98| 91| 80| 77| 84| 53| 46| 65| 44|
1627+---+---+---+---+---+---+---+---+---+---+
1628| 90| 23| 88| 95| 70| 79| 68| 83| 14| 17|
1629+---+---+---+---+---+---+---+---+---+---+
1630| 97| 92| 21| 78| 81| 94| 19| 16| 45| 66|
1631+---+---+---+---+---+---+---+---+---+---+
1632| 22| 89| 96| 93| 20| 69| 82| 67| 18| 15|
1633+---+---+---+---+---+---+---+---+---+---+
1634Solution 2
1635+---+---+---+---+---+---+---+---+---+---+
1636| 1| 58| 27| 34| 3| 40| 29| 10| 5| 8|
1637+---+---+---+---+---+---+---+---+---+---+
1638| 26| 35| 2| 57| 28| 33| 4| 7| 30| 11|
1639+---+---+---+---+---+---+---+---+---+---+
1640| 59|100| 73| 36| 41| 56| 39| 32| 9| 6|
1641+---+---+---+---+---+---+---+---+---+---+
1642| 74| 25| 60| 55| 72| 37| 42| 49| 12| 31|
1643+---+---+---+---+---+---+---+---+---+---+
1644| 61| 86| 99| 76| 63| 52| 47| 38| 43| 50|
1645+---+---+---+---+---+---+---+---+---+---+
1646| 24| 75| 62| 85| 54| 71| 64| 51| 48| 13|
1647+---+---+---+---+---+---+---+---+---+---+
1648| 87| 98| 89| 80| 77| 84| 53| 46| 65| 44|
1649+---+---+---+---+---+---+---+---+---+---+
1650| 90| 23| 92| 95| 70| 79| 68| 83| 14| 17|
1651+---+---+---+---+---+---+---+---+---+---+
1652| 97| 88| 21| 78| 81| 94| 19| 16| 45| 66|
1653+---+---+---+---+---+---+---+---+---+---+
1654| 22| 91| 96| 93| 20| 69| 82| 67| 18| 15|
1655+---+---+---+---+---+---+---+---+---+---+
Tim Petersbe4f0a72001-06-29 02:41:16 +00001656"""
1657
Fred Drake56d12662002-08-09 18:37:10 +00001658weakref_tests = """\
1659Generators are weakly referencable:
1660
1661>>> import weakref
1662>>> def gen():
1663... yield 'foo!'
1664...
1665>>> wr = weakref.ref(gen)
1666>>> wr() is gen
1667True
1668>>> p = weakref.proxy(gen)
1669
1670Generator-iterators are weakly referencable as well:
1671
1672>>> gi = gen()
1673>>> wr = weakref.ref(gi)
1674>>> wr() is gi
1675True
1676>>> p = weakref.proxy(gi)
1677>>> list(p)
1678['foo!']
1679
1680"""
1681
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001682coroutine_tests = """\
1683Sending a value into a started generator:
1684
1685>>> def f():
Guido van Rossum7131f842007-02-09 20:13:25 +00001686... print((yield 1))
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001687... yield 2
1688>>> g = f()
Georg Brandla18af4e2007-04-21 15:47:16 +00001689>>> next(g)
Phillip J. Eby0d6615f2005-08-02 00:46:46 +000016901
1691>>> g.send(42)
169242
16932
1694
1695Sending a value into a new generator produces a TypeError:
1696
1697>>> f().send("foo")
1698Traceback (most recent call last):
1699...
1700TypeError: can't send non-None value to a just-started generator
1701
1702
1703Yield by itself yields None:
1704
1705>>> def f(): yield
1706>>> list(f())
1707[None]
1708
1709
1710
1711An obscene abuse of a yield expression within a generator expression:
1712
1713>>> list((yield 21) for i in range(4))
1714[21, None, 21, None, 21, None, 21, None]
1715
1716And a more sane, but still weird usage:
1717
1718>>> def f(): list(i for i in [(yield 26)])
1719>>> type(f())
Martin v. Löwis250ad612008-04-07 05:43:42 +00001720<class 'generator'>
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001721
1722
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001723A yield expression with augmented assignment.
1724
1725>>> def coroutine(seq):
1726... count = 0
1727... while count < 200:
1728... count += yield
1729... seq.append(count)
1730>>> seq = []
1731>>> c = coroutine(seq)
Georg Brandla18af4e2007-04-21 15:47:16 +00001732>>> next(c)
Guido van Rossum7131f842007-02-09 20:13:25 +00001733>>> print(seq)
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001734[]
1735>>> c.send(10)
Guido van Rossum7131f842007-02-09 20:13:25 +00001736>>> print(seq)
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001737[10]
1738>>> c.send(10)
Guido van Rossum7131f842007-02-09 20:13:25 +00001739>>> print(seq)
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001740[10, 20]
1741>>> c.send(10)
Guido van Rossum7131f842007-02-09 20:13:25 +00001742>>> print(seq)
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001743[10, 20, 30]
1744
1745
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001746Check some syntax errors for yield expressions:
1747
1748>>> f=lambda: (yield 1),(yield 2)
1749Traceback (most recent call last):
1750 ...
Guido van Rossum33d26892007-08-05 15:29:28 +00001751SyntaxError: 'yield' outside function
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001752
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001753>>> def f(): x = yield = y
1754Traceback (most recent call last):
1755 ...
Guido van Rossum33d26892007-08-05 15:29:28 +00001756SyntaxError: assignment to yield expression not possible
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001757
1758>>> def f(): (yield bar) = y
1759Traceback (most recent call last):
1760 ...
Guido van Rossum33d26892007-08-05 15:29:28 +00001761SyntaxError: can't assign to yield expression
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001762
1763>>> def f(): (yield bar) += y
1764Traceback (most recent call last):
1765 ...
Benjamin Peterson87c8d872009-06-11 22:54:11 +00001766SyntaxError: can't assign to yield expression
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001767
1768
1769Now check some throw() conditions:
1770
1771>>> def f():
1772... while True:
1773... try:
Guido van Rossum7131f842007-02-09 20:13:25 +00001774... print((yield))
Guido van Rossumb940e112007-01-10 16:19:56 +00001775... except ValueError as v:
Guido van Rossumc420b2f2007-02-09 22:09:01 +00001776... print("caught ValueError (%s)" % (v))
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001777>>> import sys
1778>>> g = f()
Georg Brandla18af4e2007-04-21 15:47:16 +00001779>>> next(g)
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001780
1781>>> g.throw(ValueError) # type only
1782caught ValueError ()
1783
1784>>> g.throw(ValueError("xyz")) # value only
1785caught ValueError (xyz)
1786
1787>>> g.throw(ValueError, ValueError(1)) # value+matching type
1788caught ValueError (1)
1789
1790>>> g.throw(ValueError, TypeError(1)) # mismatched type, rewrapped
1791caught ValueError (1)
1792
Thomas Wouters49fd7fa2006-04-21 10:40:58 +00001793>>> g.throw(ValueError, ValueError(1), None) # explicit None traceback
1794caught ValueError (1)
1795
Tim Peterse9fe7e02005-08-07 03:04:58 +00001796>>> g.throw(ValueError(1), "foo") # bad args
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001797Traceback (most recent call last):
1798 ...
1799TypeError: instance exception may not have a separate value
1800
Tim Peterse9fe7e02005-08-07 03:04:58 +00001801>>> g.throw(ValueError, "foo", 23) # bad args
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001802Traceback (most recent call last):
1803 ...
1804TypeError: throw() third argument must be a traceback object
1805
Guido van Rossumbf12cdb2006-08-17 20:24:18 +00001806>>> g.throw("abc")
1807Traceback (most recent call last):
1808 ...
1809TypeError: exceptions must be classes or instances deriving from BaseException, not str
1810
1811>>> g.throw(0)
1812Traceback (most recent call last):
1813 ...
1814TypeError: exceptions must be classes or instances deriving from BaseException, not int
1815
1816>>> g.throw(list)
1817Traceback (most recent call last):
1818 ...
1819TypeError: exceptions must be classes or instances deriving from BaseException, not type
1820
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001821>>> def throw(g,exc):
1822... try:
1823... raise exc
1824... except:
1825... g.throw(*sys.exc_info())
Tim Peterse9fe7e02005-08-07 03:04:58 +00001826>>> throw(g,ValueError) # do it with traceback included
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001827caught ValueError ()
1828
1829>>> g.send(1)
18301
1831
Tim Peterse9fe7e02005-08-07 03:04:58 +00001832>>> throw(g,TypeError) # terminate the generator
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001833Traceback (most recent call last):
1834 ...
1835TypeError
1836
Guido van Rossum7131f842007-02-09 20:13:25 +00001837>>> print(g.gi_frame)
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001838None
1839
1840>>> g.send(2)
1841Traceback (most recent call last):
1842 ...
1843StopIteration
1844
Tim Peterse9fe7e02005-08-07 03:04:58 +00001845>>> g.throw(ValueError,6) # throw on closed generator
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001846Traceback (most recent call last):
1847 ...
1848ValueError: 6
1849
Tim Peterse9fe7e02005-08-07 03:04:58 +00001850>>> f().throw(ValueError,7) # throw on just-opened generator
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001851Traceback (most recent call last):
1852 ...
1853ValueError: 7
1854
Antoine Pitrou551ba202011-10-18 16:40:50 +02001855Plain "raise" inside a generator should preserve the traceback (#13188).
1856The traceback should have 3 levels:
1857- g.throw()
1858- f()
1859- 1/0
1860
1861>>> def f():
1862... try:
1863... yield
1864... except:
1865... raise
1866>>> g = f()
1867>>> try:
1868... 1/0
1869... except ZeroDivisionError as v:
1870... try:
1871... g.throw(v)
1872... except Exception as w:
1873... tb = w.__traceback__
1874>>> levels = 0
1875>>> while tb:
1876... levels += 1
1877... tb = tb.tb_next
1878>>> levels
18793
1880
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001881Now let's try closing a generator:
1882
1883>>> def f():
1884... try: yield
1885... except GeneratorExit:
Guido van Rossum7131f842007-02-09 20:13:25 +00001886... print("exiting")
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001887
1888>>> g = f()
Georg Brandla18af4e2007-04-21 15:47:16 +00001889>>> next(g)
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001890>>> g.close()
1891exiting
1892>>> g.close() # should be no-op now
1893
1894>>> f().close() # close on just-opened generator should be fine
1895
Tim Peterse9fe7e02005-08-07 03:04:58 +00001896>>> def f(): yield # an even simpler generator
1897>>> f().close() # close before opening
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001898>>> g = f()
Georg Brandla18af4e2007-04-21 15:47:16 +00001899>>> next(g)
Tim Peterse9fe7e02005-08-07 03:04:58 +00001900>>> g.close() # close normally
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001901
1902And finalization:
1903
1904>>> def f():
1905... try: yield
1906... finally:
Guido van Rossum7131f842007-02-09 20:13:25 +00001907... print("exiting")
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001908
1909>>> g = f()
Georg Brandla18af4e2007-04-21 15:47:16 +00001910>>> next(g)
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001911>>> del g
1912exiting
1913
1914
Christian Heimescbf3b5c2007-12-03 21:02:03 +00001915GeneratorExit is not caught by except Exception:
1916
1917>>> def f():
1918... try: yield
1919... except Exception:
1920... print('except')
1921... finally:
1922... print('finally')
1923
1924>>> g = f()
1925>>> next(g)
1926>>> del g
1927finally
1928
1929
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001930Now let's try some ill-behaved generators:
1931
1932>>> def f():
1933... try: yield
1934... except GeneratorExit:
1935... yield "foo!"
1936>>> g = f()
Georg Brandla18af4e2007-04-21 15:47:16 +00001937>>> next(g)
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001938>>> g.close()
1939Traceback (most recent call last):
1940 ...
1941RuntimeError: generator ignored GeneratorExit
1942>>> g.close()
1943
1944
1945Our ill-behaved code should be invoked during GC:
1946
Guido van Rossum34d19282007-08-09 01:03:29 +00001947>>> import sys, io
1948>>> old, sys.stderr = sys.stderr, io.StringIO()
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001949>>> g = f()
Georg Brandla18af4e2007-04-21 15:47:16 +00001950>>> next(g)
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001951>>> del g
Andrew Svetlov76bcff22012-11-03 15:56:05 +02001952>>> "RuntimeError: generator ignored GeneratorExit" in sys.stderr.getvalue()
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001953True
1954>>> sys.stderr = old
1955
1956
1957And errors thrown during closing should propagate:
1958
1959>>> def f():
1960... try: yield
1961... except GeneratorExit:
1962... raise TypeError("fie!")
1963>>> g = f()
Georg Brandla18af4e2007-04-21 15:47:16 +00001964>>> next(g)
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001965>>> g.close()
1966Traceback (most recent call last):
1967 ...
1968TypeError: fie!
1969
1970
1971Ensure that various yield expression constructs make their
1972enclosing function a generator:
1973
1974>>> def f(): x += yield
1975>>> type(f())
Martin v. Löwis250ad612008-04-07 05:43:42 +00001976<class 'generator'>
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001977
1978>>> def f(): x = yield
1979>>> type(f())
Martin v. Löwis250ad612008-04-07 05:43:42 +00001980<class 'generator'>
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001981
1982>>> def f(): lambda x=(yield): 1
1983>>> type(f())
Martin v. Löwis250ad612008-04-07 05:43:42 +00001984<class 'generator'>
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001985
1986>>> def f(): x=(i for i in (yield) if (yield))
1987>>> type(f())
Martin v. Löwis250ad612008-04-07 05:43:42 +00001988<class 'generator'>
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001989
1990>>> def f(d): d[(yield "a")] = d[(yield "b")] = 27
1991>>> data = [1,2]
1992>>> g = f(data)
1993>>> type(g)
Martin v. Löwis250ad612008-04-07 05:43:42 +00001994<class 'generator'>
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001995>>> g.send(None)
1996'a'
1997>>> data
1998[1, 2]
1999>>> g.send(0)
2000'b'
2001>>> data
2002[27, 2]
2003>>> try: g.send(1)
2004... except StopIteration: pass
2005>>> data
2006[27, 27]
2007
2008"""
2009
Thomas Wouters49fd7fa2006-04-21 10:40:58 +00002010refleaks_tests = """
2011Prior to adding cycle-GC support to itertools.tee, this code would leak
2012references. We add it to the standard suite so the routine refleak-tests
2013would trigger if it starts being uncleanable again.
2014
2015>>> import itertools
2016>>> def leak():
2017... class gen:
2018... def __iter__(self):
2019... return self
Georg Brandla18af4e2007-04-21 15:47:16 +00002020... def __next__(self):
Thomas Wouters49fd7fa2006-04-21 10:40:58 +00002021... return self.item
2022... g = gen()
2023... head, tail = itertools.tee(g)
2024... g.item = head
2025... return head
2026>>> it = leak()
2027
2028Make sure to also test the involvement of the tee-internal teedataobject,
2029which stores returned items.
2030
Georg Brandla18af4e2007-04-21 15:47:16 +00002031>>> item = next(it)
Thomas Wouters49fd7fa2006-04-21 10:40:58 +00002032
2033
2034
2035This test leaked at one point due to generator finalization/destruction.
2036It was copied from Lib/test/leakers/test_generator_cycle.py before the file
2037was removed.
2038
2039>>> def leak():
2040... def gen():
2041... while True:
2042... yield g
2043... g = gen()
2044
2045>>> leak()
2046
2047
2048
2049This test isn't really generator related, but rather exception-in-cleanup
2050related. The coroutine tests (above) just happen to cause an exception in
2051the generator's __del__ (tp_del) method. We can also test for this
2052explicitly, without generators. We do have to redirect stderr to avoid
2053printing warnings and to doublecheck that we actually tested what we wanted
2054to test.
2055
Guido van Rossum34d19282007-08-09 01:03:29 +00002056>>> import sys, io
Thomas Wouters49fd7fa2006-04-21 10:40:58 +00002057>>> old = sys.stderr
2058>>> try:
Guido van Rossum34d19282007-08-09 01:03:29 +00002059... sys.stderr = io.StringIO()
Thomas Wouters49fd7fa2006-04-21 10:40:58 +00002060... class Leaker:
2061... def __del__(self):
Andrew Svetlov76bcff22012-11-03 15:56:05 +02002062... def invoke(message):
2063... raise RuntimeError(message)
2064... invoke("test")
Thomas Wouters49fd7fa2006-04-21 10:40:58 +00002065...
2066... l = Leaker()
2067... del l
2068... err = sys.stderr.getvalue().strip()
Andrew Svetlov76bcff22012-11-03 15:56:05 +02002069... "Exception ignored in" in err
2070... "RuntimeError: test" in err
2071... "Traceback" in err
2072... "in invoke" in err
Thomas Wouters49fd7fa2006-04-21 10:40:58 +00002073... finally:
2074... sys.stderr = old
2075True
2076True
Andrew Svetlov76bcff22012-11-03 15:56:05 +02002077True
2078True
Thomas Wouters49fd7fa2006-04-21 10:40:58 +00002079
2080
2081These refleak tests should perhaps be in a testfile of their own,
2082test_generators just happened to be the test that drew these out.
2083
2084"""
2085
Tim Petersf6ed0742001-06-27 07:17:57 +00002086__test__ = {"tut": tutorial_tests,
2087 "pep": pep_tests,
2088 "email": email_tests,
2089 "fun": fun_tests,
Tim Petersbe4f0a72001-06-29 02:41:16 +00002090 "syntax": syntax_tests,
Fred Drake56d12662002-08-09 18:37:10 +00002091 "conjoin": conjoin_tests,
2092 "weakref": weakref_tests,
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00002093 "coroutine": coroutine_tests,
Thomas Wouters49fd7fa2006-04-21 10:40:58 +00002094 "refleaks": refleaks_tests,
Fred Drake56d12662002-08-09 18:37:10 +00002095 }
Tim Peters1def3512001-06-23 20:27:04 +00002096
2097# Magic test name that regrtest.py invokes *after* importing this module.
2098# This worms around a bootstrap problem.
2099# Note that doctest and regrtest both look in sys.argv for a "-v" argument,
2100# so this works as expected in both ways of running regrtest.
Tim Petersa0a62222001-09-09 06:12:01 +00002101def test_main(verbose=None):
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002102 from test import support, test_generators
Antoine Pitrou796564c2013-07-30 19:59:21 +02002103 support.run_unittest(__name__)
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002104 support.run_doctest(test_generators, verbose)
Tim Peters1def3512001-06-23 20:27:04 +00002105
2106# This part isn't needed for regrtest, but for running the test directly.
2107if __name__ == "__main__":
Tim Petersa0a62222001-09-09 06:12:01 +00002108 test_main(1)