blob: 604e12d92e158dbce5d00fecc76a9542b5db9fb0 [file] [log] [blame]
Serhiy Storchakad7a44152015-11-12 11:23:04 +02001import copy
Antoine Pitrou796564c2013-07-30 19:59:21 +02002import gc
Serhiy Storchakad7a44152015-11-12 11:23:04 +02003import pickle
Antoine Pitrou796564c2013-07-30 19:59:21 +02004import sys
5import unittest
6import weakref
7
8from test import support
9
10
11class FinalizationTest(unittest.TestCase):
12
13 def test_frame_resurrect(self):
14 # A generator frame can be resurrected by a generator's finalization.
15 def gen():
16 nonlocal frame
17 try:
18 yield
19 finally:
20 frame = sys._getframe()
21
22 g = gen()
23 wr = weakref.ref(g)
24 next(g)
25 del g
26 support.gc_collect()
27 self.assertIs(wr(), None)
28 self.assertTrue(frame)
29 del frame
30 support.gc_collect()
31
32 def test_refcycle(self):
33 # A generator caught in a refcycle gets finalized anyway.
34 old_garbage = gc.garbage[:]
35 finalized = False
36 def gen():
37 nonlocal finalized
38 try:
39 g = yield
40 yield 1
41 finally:
42 finalized = True
43
44 g = gen()
45 next(g)
46 g.send(g)
47 self.assertGreater(sys.getrefcount(g), 2)
48 self.assertFalse(finalized)
49 del g
50 support.gc_collect()
51 self.assertTrue(finalized)
52 self.assertEqual(gc.garbage, old_garbage)
53
Serhiy Storchakac775ad62015-03-11 18:20:35 +020054 def test_lambda_generator(self):
55 # Issue #23192: Test that a lambda returning a generator behaves
56 # like the equivalent function
57 f = lambda: (yield 1)
58 def g(): return (yield 1)
59
60 # test 'yield from'
61 f2 = lambda: (yield from g())
62 def g2(): return (yield from g())
63
64 f3 = lambda: (yield from f())
65 def g3(): return (yield from f())
66
67 for gen_fun in (f, g, f2, g2, f3, g3):
68 gen = gen_fun()
69 self.assertEqual(next(gen), 1)
70 with self.assertRaises(StopIteration) as cm:
71 gen.send(2)
72 self.assertEqual(cm.exception.value, 2)
73
Antoine Pitrou796564c2013-07-30 19:59:21 +020074
Serhiy Storchakad7a44152015-11-12 11:23:04 +020075class GeneratorTest(unittest.TestCase):
76
77 def test_copy(self):
78 def f():
79 yield 1
80 g = f()
81 with self.assertRaises(TypeError):
82 copy.copy(g)
83
84 def test_pickle(self):
85 def f():
86 yield 1
87 g = f()
88 for proto in range(pickle.HIGHEST_PROTOCOL + 1):
89 with self.assertRaises((TypeError, pickle.PicklingError)):
90 pickle.dumps(g, proto)
91
92
Victor Stinner26f7b8a2015-01-31 10:29:47 +010093class ExceptionTest(unittest.TestCase):
94 # Tests for the issue #23353: check that the currently handled exception
95 # is correctly saved/restored in PyEval_EvalFrameEx().
96
97 def test_except_throw(self):
98 def store_raise_exc_generator():
99 try:
100 self.assertEqual(sys.exc_info()[0], None)
101 yield
102 except Exception as exc:
103 # exception raised by gen.throw(exc)
104 self.assertEqual(sys.exc_info()[0], ValueError)
105 self.assertIsNone(exc.__context__)
106 yield
107
108 # ensure that the exception is not lost
109 self.assertEqual(sys.exc_info()[0], ValueError)
110 yield
111
112 # we should be able to raise back the ValueError
113 raise
114
115 make = store_raise_exc_generator()
116 next(make)
117
118 try:
119 raise ValueError()
120 except Exception as exc:
121 try:
122 make.throw(exc)
123 except Exception:
124 pass
125
126 next(make)
127 with self.assertRaises(ValueError) as cm:
128 next(make)
129 self.assertIsNone(cm.exception.__context__)
130
131 self.assertEqual(sys.exc_info(), (None, None, None))
132
133 def test_except_next(self):
134 def gen():
135 self.assertEqual(sys.exc_info()[0], ValueError)
136 yield "done"
137
138 g = gen()
139 try:
140 raise ValueError
141 except Exception:
142 self.assertEqual(next(g), "done")
143 self.assertEqual(sys.exc_info(), (None, None, None))
144
145 def test_except_gen_except(self):
146 def gen():
147 try:
148 self.assertEqual(sys.exc_info()[0], None)
149 yield
150 # we are called from "except ValueError:", TypeError must
151 # inherit ValueError in its context
152 raise TypeError()
153 except TypeError as exc:
154 self.assertEqual(sys.exc_info()[0], TypeError)
155 self.assertEqual(type(exc.__context__), ValueError)
156 # here we are still called from the "except ValueError:"
157 self.assertEqual(sys.exc_info()[0], ValueError)
158 yield
159 self.assertIsNone(sys.exc_info()[0])
160 yield "done"
161
162 g = gen()
163 next(g)
164 try:
165 raise ValueError
166 except Exception:
167 next(g)
168
169 self.assertEqual(next(g), "done")
170 self.assertEqual(sys.exc_info(), (None, None, None))
171
172 def test_except_throw_exception_context(self):
173 def gen():
174 try:
175 try:
176 self.assertEqual(sys.exc_info()[0], None)
177 yield
178 except ValueError:
179 # we are called from "except ValueError:"
180 self.assertEqual(sys.exc_info()[0], ValueError)
181 raise TypeError()
182 except Exception as exc:
183 self.assertEqual(sys.exc_info()[0], TypeError)
184 self.assertEqual(type(exc.__context__), ValueError)
185 # we are still called from "except ValueError:"
186 self.assertEqual(sys.exc_info()[0], ValueError)
187 yield
188 self.assertIsNone(sys.exc_info()[0])
189 yield "done"
190
191 g = gen()
192 next(g)
193 try:
194 raise ValueError
195 except Exception as exc:
196 g.throw(exc)
197
198 self.assertEqual(next(g), "done")
199 self.assertEqual(sys.exc_info(), (None, None, None))
200
201
Tim Peters6ba5f792001-06-23 20:45:43 +0000202tutorial_tests = """
Tim Peters1def3512001-06-23 20:27:04 +0000203Let's try a simple generator:
204
205 >>> def f():
206 ... yield 1
207 ... yield 2
208
Tim Petersb9e9ff12001-06-24 03:44:52 +0000209 >>> for i in f():
Guido van Rossum7131f842007-02-09 20:13:25 +0000210 ... print(i)
Tim Petersb9e9ff12001-06-24 03:44:52 +0000211 1
212 2
Tim Peters1def3512001-06-23 20:27:04 +0000213 >>> g = f()
Georg Brandla18af4e2007-04-21 15:47:16 +0000214 >>> next(g)
Tim Peters1def3512001-06-23 20:27:04 +0000215 1
Georg Brandla18af4e2007-04-21 15:47:16 +0000216 >>> next(g)
Tim Peters1def3512001-06-23 20:27:04 +0000217 2
Tim Petersea2e97a2001-06-24 07:10:02 +0000218
Tim Peters2106ef02001-06-25 01:30:12 +0000219"Falling off the end" stops the generator:
Tim Petersea2e97a2001-06-24 07:10:02 +0000220
Georg Brandla18af4e2007-04-21 15:47:16 +0000221 >>> next(g)
Tim Peters1def3512001-06-23 20:27:04 +0000222 Traceback (most recent call last):
223 File "<stdin>", line 1, in ?
224 File "<stdin>", line 2, in g
225 StopIteration
226
Tim Petersea2e97a2001-06-24 07:10:02 +0000227"return" also stops the generator:
Tim Peters1def3512001-06-23 20:27:04 +0000228
229 >>> def f():
230 ... yield 1
231 ... return
232 ... yield 2 # never reached
233 ...
234 >>> g = f()
Georg Brandla18af4e2007-04-21 15:47:16 +0000235 >>> next(g)
Tim Peters1def3512001-06-23 20:27:04 +0000236 1
Georg Brandla18af4e2007-04-21 15:47:16 +0000237 >>> next(g)
Tim Peters1def3512001-06-23 20:27:04 +0000238 Traceback (most recent call last):
239 File "<stdin>", line 1, in ?
240 File "<stdin>", line 3, in f
241 StopIteration
Georg Brandla18af4e2007-04-21 15:47:16 +0000242 >>> next(g) # once stopped, can't be resumed
Tim Peters1def3512001-06-23 20:27:04 +0000243 Traceback (most recent call last):
244 File "<stdin>", line 1, in ?
245 StopIteration
246
247"raise StopIteration" stops the generator too:
248
249 >>> def f():
250 ... yield 1
Tim Peters34463652001-07-12 22:43:41 +0000251 ... raise StopIteration
Tim Peters1def3512001-06-23 20:27:04 +0000252 ... yield 2 # never reached
253 ...
254 >>> g = f()
Georg Brandla18af4e2007-04-21 15:47:16 +0000255 >>> next(g)
Tim Peters1def3512001-06-23 20:27:04 +0000256 1
Georg Brandla18af4e2007-04-21 15:47:16 +0000257 >>> next(g)
Tim Peters1def3512001-06-23 20:27:04 +0000258 Traceback (most recent call last):
259 File "<stdin>", line 1, in ?
260 StopIteration
Georg Brandla18af4e2007-04-21 15:47:16 +0000261 >>> next(g)
Tim Peters1def3512001-06-23 20:27:04 +0000262 Traceback (most recent call last):
263 File "<stdin>", line 1, in ?
264 StopIteration
265
266However, they are not exactly equivalent:
267
268 >>> def g1():
269 ... try:
270 ... return
271 ... except:
272 ... yield 1
273 ...
274 >>> list(g1())
275 []
276
277 >>> def g2():
278 ... try:
279 ... raise StopIteration
280 ... except:
281 ... yield 42
Guido van Rossum7131f842007-02-09 20:13:25 +0000282 >>> print(list(g2()))
Tim Peters1def3512001-06-23 20:27:04 +0000283 [42]
284
285This may be surprising at first:
286
287 >>> def g3():
288 ... try:
289 ... return
290 ... finally:
291 ... yield 1
292 ...
293 >>> list(g3())
294 [1]
295
296Let's create an alternate range() function implemented as a generator:
297
298 >>> def yrange(n):
299 ... for i in range(n):
300 ... yield i
301 ...
302 >>> list(yrange(5))
303 [0, 1, 2, 3, 4]
304
305Generators always return to the most recent caller:
306
307 >>> def creator():
308 ... r = yrange(5)
Georg Brandla18af4e2007-04-21 15:47:16 +0000309 ... print("creator", next(r))
Tim Peters1def3512001-06-23 20:27:04 +0000310 ... return r
Guido van Rossumc420b2f2007-02-09 22:09:01 +0000311 ...
Tim Peters1def3512001-06-23 20:27:04 +0000312 >>> def caller():
313 ... r = creator()
314 ... for i in r:
Guido van Rossum7131f842007-02-09 20:13:25 +0000315 ... print("caller", i)
Guido van Rossumc420b2f2007-02-09 22:09:01 +0000316 ...
Tim Peters1def3512001-06-23 20:27:04 +0000317 >>> caller()
318 creator 0
319 caller 1
320 caller 2
321 caller 3
322 caller 4
323
324Generators can call other generators:
325
326 >>> def zrange(n):
327 ... for i in yrange(n):
328 ... yield i
329 ...
330 >>> list(zrange(5))
331 [0, 1, 2, 3, 4]
332
333"""
334
Tim Peters6ba5f792001-06-23 20:45:43 +0000335# The examples from PEP 255.
336
337pep_tests = """
338
Tim Peterse5614632001-08-15 04:41:19 +0000339Specification: Yield
340
341 Restriction: A generator cannot be resumed while it is actively
342 running:
343
344 >>> def g():
Georg Brandla18af4e2007-04-21 15:47:16 +0000345 ... i = next(me)
Tim Peterse5614632001-08-15 04:41:19 +0000346 ... yield i
347 >>> me = g()
Georg Brandla18af4e2007-04-21 15:47:16 +0000348 >>> next(me)
Tim Peterse5614632001-08-15 04:41:19 +0000349 Traceback (most recent call last):
350 ...
351 File "<string>", line 2, in g
352 ValueError: generator already executing
353
Tim Peters6ba5f792001-06-23 20:45:43 +0000354Specification: Return
355
356 Note that return isn't always equivalent to raising StopIteration: the
357 difference lies in how enclosing try/except constructs are treated.
358 For example,
359
360 >>> def f1():
361 ... try:
362 ... return
363 ... except:
364 ... yield 1
Guido van Rossum7131f842007-02-09 20:13:25 +0000365 >>> print(list(f1()))
Tim Peters6ba5f792001-06-23 20:45:43 +0000366 []
367
368 because, as in any function, return simply exits, but
369
370 >>> def f2():
371 ... try:
372 ... raise StopIteration
373 ... except:
374 ... yield 42
Guido van Rossum7131f842007-02-09 20:13:25 +0000375 >>> print(list(f2()))
Tim Peters6ba5f792001-06-23 20:45:43 +0000376 [42]
377
378 because StopIteration is captured by a bare "except", as is any
379 exception.
380
381Specification: Generators and Exception Propagation
382
383 >>> def f():
Tim Peters3caca232001-12-06 06:23:26 +0000384 ... return 1//0
Tim Peters6ba5f792001-06-23 20:45:43 +0000385 >>> def g():
386 ... yield f() # the zero division exception propagates
387 ... yield 42 # and we'll never get here
388 >>> k = g()
Georg Brandla18af4e2007-04-21 15:47:16 +0000389 >>> next(k)
Tim Peters6ba5f792001-06-23 20:45:43 +0000390 Traceback (most recent call last):
391 File "<stdin>", line 1, in ?
392 File "<stdin>", line 2, in g
393 File "<stdin>", line 2, in f
394 ZeroDivisionError: integer division or modulo by zero
Georg Brandla18af4e2007-04-21 15:47:16 +0000395 >>> next(k) # and the generator cannot be resumed
Tim Peters6ba5f792001-06-23 20:45:43 +0000396 Traceback (most recent call last):
397 File "<stdin>", line 1, in ?
398 StopIteration
399 >>>
400
401Specification: Try/Except/Finally
402
403 >>> def f():
404 ... try:
405 ... yield 1
406 ... try:
407 ... yield 2
Tim Peters3caca232001-12-06 06:23:26 +0000408 ... 1//0
Tim Peters6ba5f792001-06-23 20:45:43 +0000409 ... yield 3 # never get here
410 ... except ZeroDivisionError:
411 ... yield 4
412 ... yield 5
413 ... raise
414 ... except:
415 ... yield 6
416 ... yield 7 # the "raise" above stops this
417 ... except:
418 ... yield 8
419 ... yield 9
420 ... try:
421 ... x = 12
422 ... finally:
423 ... yield 10
424 ... yield 11
Guido van Rossum7131f842007-02-09 20:13:25 +0000425 >>> print(list(f()))
Tim Peters6ba5f792001-06-23 20:45:43 +0000426 [1, 2, 4, 5, 8, 9, 10, 11]
427 >>>
428
Tim Peters6ba5f792001-06-23 20:45:43 +0000429Guido's binary tree example.
430
431 >>> # A binary tree class.
432 >>> class Tree:
433 ...
434 ... def __init__(self, label, left=None, right=None):
435 ... self.label = label
436 ... self.left = left
437 ... self.right = right
438 ...
439 ... def __repr__(self, level=0, indent=" "):
Walter Dörwald70a6b492004-02-12 17:35:32 +0000440 ... s = level*indent + repr(self.label)
Tim Peters6ba5f792001-06-23 20:45:43 +0000441 ... if self.left:
442 ... s = s + "\\n" + self.left.__repr__(level+1, indent)
443 ... if self.right:
444 ... s = s + "\\n" + self.right.__repr__(level+1, indent)
445 ... return s
446 ...
447 ... def __iter__(self):
448 ... return inorder(self)
449
450 >>> # Create a Tree from a list.
451 >>> def tree(list):
452 ... n = len(list)
453 ... if n == 0:
454 ... return []
Tim Peters3caca232001-12-06 06:23:26 +0000455 ... i = n // 2
Tim Peters6ba5f792001-06-23 20:45:43 +0000456 ... return Tree(list[i], tree(list[:i]), tree(list[i+1:]))
457
458 >>> # Show it off: create a tree.
459 >>> t = tree("ABCDEFGHIJKLMNOPQRSTUVWXYZ")
460
Tim Petersd674e172002-03-10 07:59:13 +0000461 >>> # A recursive generator that generates Tree labels in in-order.
Tim Peters6ba5f792001-06-23 20:45:43 +0000462 >>> def inorder(t):
463 ... if t:
464 ... for x in inorder(t.left):
465 ... yield x
466 ... yield t.label
467 ... for x in inorder(t.right):
468 ... yield x
469
470 >>> # Show it off: create a tree.
Edward Loper103d26e2004-08-09 02:03:30 +0000471 >>> t = tree("ABCDEFGHIJKLMNOPQRSTUVWXYZ")
472 >>> # Print the nodes of the tree in in-order.
473 >>> for x in t:
Guido van Rossumc420b2f2007-02-09 22:09:01 +0000474 ... print(' '+x, end='')
475 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 +0000476
477 >>> # A non-recursive generator.
478 >>> def inorder(node):
479 ... stack = []
480 ... while node:
481 ... while node.left:
482 ... stack.append(node)
483 ... node = node.left
484 ... yield node.label
485 ... while not node.right:
486 ... try:
487 ... node = stack.pop()
488 ... except IndexError:
489 ... return
490 ... yield node.label
491 ... node = node.right
492
493 >>> # Exercise the non-recursive generator.
494 >>> for x in t:
Guido van Rossumc420b2f2007-02-09 22:09:01 +0000495 ... print(' '+x, end='')
496 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 +0000497
498"""
499
Tim Petersb2bc6a92001-06-24 10:14:27 +0000500# Examples from Iterator-List and Python-Dev and c.l.py.
Tim Peters6ba5f792001-06-23 20:45:43 +0000501
502email_tests = """
503
504The difference between yielding None and returning it.
505
506>>> def g():
507... for i in range(3):
508... yield None
509... yield None
510... return
511>>> list(g())
512[None, None, None, None]
513
514Ensure that explicitly raising StopIteration acts like any other exception
515in try/except, not like a return.
516
517>>> def g():
518... yield 1
519... try:
520... raise StopIteration
521... except:
522... yield 2
523... yield 3
524>>> list(g())
525[1, 2, 3]
Tim Petersb9e9ff12001-06-24 03:44:52 +0000526
Tim Petersb2bc6a92001-06-24 10:14:27 +0000527Next one was posted to c.l.py.
528
529>>> def gcomb(x, k):
530... "Generate all combinations of k elements from list x."
531...
532... if k > len(x):
533... return
534... if k == 0:
535... yield []
536... else:
537... first, rest = x[0], x[1:]
538... # A combination does or doesn't contain first.
539... # If it does, the remainder is a k-1 comb of rest.
540... for c in gcomb(rest, k-1):
541... c.insert(0, first)
542... yield c
543... # If it doesn't contain first, it's a k comb of rest.
544... for c in gcomb(rest, k):
545... yield c
546
Guido van Rossum805365e2007-05-07 22:24:25 +0000547>>> seq = list(range(1, 5))
Tim Petersb2bc6a92001-06-24 10:14:27 +0000548>>> for k in range(len(seq) + 2):
Guido van Rossum7131f842007-02-09 20:13:25 +0000549... print("%d-combs of %s:" % (k, seq))
Tim Petersb2bc6a92001-06-24 10:14:27 +0000550... for c in gcomb(seq, k):
Guido van Rossum7131f842007-02-09 20:13:25 +0000551... print(" ", c)
Tim Petersb2bc6a92001-06-24 10:14:27 +00005520-combs of [1, 2, 3, 4]:
553 []
5541-combs of [1, 2, 3, 4]:
555 [1]
556 [2]
557 [3]
558 [4]
5592-combs of [1, 2, 3, 4]:
560 [1, 2]
561 [1, 3]
562 [1, 4]
563 [2, 3]
564 [2, 4]
565 [3, 4]
5663-combs of [1, 2, 3, 4]:
567 [1, 2, 3]
568 [1, 2, 4]
569 [1, 3, 4]
570 [2, 3, 4]
5714-combs of [1, 2, 3, 4]:
572 [1, 2, 3, 4]
5735-combs of [1, 2, 3, 4]:
Tim Peters3e7b1a02001-06-25 19:46:25 +0000574
Tim Peterse77f2e22001-06-26 22:24:51 +0000575From the Iterators list, about the types of these things.
Tim Peters3e7b1a02001-06-25 19:46:25 +0000576
577>>> def g():
578... yield 1
579...
580>>> type(g)
Martin v. Löwis250ad612008-04-07 05:43:42 +0000581<class 'function'>
Tim Peters3e7b1a02001-06-25 19:46:25 +0000582>>> i = g()
583>>> type(i)
Martin v. Löwis250ad612008-04-07 05:43:42 +0000584<class 'generator'>
Tim Peters5d2b77c2001-09-03 05:47:38 +0000585>>> [s for s in dir(i) if not s.startswith('_')]
Christian Heimesaf98da12008-01-27 15:18:18 +0000586['close', 'gi_code', 'gi_frame', 'gi_running', 'send', 'throw']
Serhiy Storchaka9a11f172013-01-31 16:11:04 +0200587>>> from test.support import HAVE_DOCSTRINGS
Larry Hastings581ee362014-01-28 05:00:08 -0800588>>> print(i.__next__.__doc__ if HAVE_DOCSTRINGS else 'Implement next(self).')
589Implement next(self).
Tim Peters3e7b1a02001-06-25 19:46:25 +0000590>>> iter(i) is i
Guido van Rossum77f6a652002-04-03 22:41:51 +0000591True
Tim Peters3e7b1a02001-06-25 19:46:25 +0000592>>> import types
593>>> isinstance(i, types.GeneratorType)
Guido van Rossum77f6a652002-04-03 22:41:51 +0000594True
Tim Peterse77f2e22001-06-26 22:24:51 +0000595
596And more, added later.
597
598>>> i.gi_running
5990
600>>> type(i.gi_frame)
Martin v. Löwis250ad612008-04-07 05:43:42 +0000601<class 'frame'>
Tim Peterse77f2e22001-06-26 22:24:51 +0000602>>> i.gi_running = 42
603Traceback (most recent call last):
604 ...
Collin Winter42dae6a2007-03-28 21:44:53 +0000605AttributeError: readonly attribute
Tim Peterse77f2e22001-06-26 22:24:51 +0000606>>> def g():
607... yield me.gi_running
608>>> me = g()
609>>> me.gi_running
6100
Georg Brandla18af4e2007-04-21 15:47:16 +0000611>>> next(me)
Tim Peterse77f2e22001-06-26 22:24:51 +00006121
613>>> me.gi_running
6140
Tim Peters35302662001-07-02 01:38:33 +0000615
616A clever union-find implementation from c.l.py, due to David Eppstein.
617Sent: Friday, June 29, 2001 12:16 PM
618To: python-list@python.org
619Subject: Re: PEP 255: Simple Generators
620
621>>> class disjointSet:
622... def __init__(self, name):
623... self.name = name
624... self.parent = None
625... self.generator = self.generate()
626...
627... def generate(self):
628... while not self.parent:
629... yield self
630... for x in self.parent.generator:
631... yield x
632...
633... def find(self):
Georg Brandla18af4e2007-04-21 15:47:16 +0000634... return next(self.generator)
Tim Peters35302662001-07-02 01:38:33 +0000635...
636... def union(self, parent):
637... if self.parent:
638... raise ValueError("Sorry, I'm not a root!")
639... self.parent = parent
640...
641... def __str__(self):
642... return self.name
Tim Peters35302662001-07-02 01:38:33 +0000643
644>>> names = "ABCDEFGHIJKLM"
645>>> sets = [disjointSet(name) for name in names]
646>>> roots = sets[:]
647
648>>> import random
Raymond Hettinger28de64f2008-01-13 23:40:30 +0000649>>> gen = random.Random(42)
Tim Peters35302662001-07-02 01:38:33 +0000650>>> while 1:
651... for s in sets:
Guido van Rossumc420b2f2007-02-09 22:09:01 +0000652... print(" %s->%s" % (s, s.find()), end='')
Guido van Rossum7131f842007-02-09 20:13:25 +0000653... print()
Tim Peters35302662001-07-02 01:38:33 +0000654... if len(roots) > 1:
Raymond Hettingerdd24a9f2002-12-30 00:46:09 +0000655... s1 = gen.choice(roots)
Tim Peters35302662001-07-02 01:38:33 +0000656... roots.remove(s1)
Raymond Hettingerdd24a9f2002-12-30 00:46:09 +0000657... s2 = gen.choice(roots)
Tim Peters35302662001-07-02 01:38:33 +0000658... s1.union(s2)
Guido van Rossum7131f842007-02-09 20:13:25 +0000659... print("merged", s1, "into", s2)
Tim Peters35302662001-07-02 01:38:33 +0000660... else:
661... break
Guido van Rossumc420b2f2007-02-09 22:09:01 +0000662 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 +0000663merged K into B
664 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 +0000665merged A into F
Raymond Hettingerc585eec2010-09-07 15:00:15 +0000666 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
667merged E into F
668 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
669merged D into C
670 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 +0000671merged M into C
Raymond Hettingerc585eec2010-09-07 15:00:15 +0000672 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
673merged J into B
674 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
675merged B into C
676 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
677merged F into G
678 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
679merged L into C
680 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
681merged G into I
682 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
683merged I into H
684 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
685merged C into H
686 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 +0000687
Tim Peters6ba5f792001-06-23 20:45:43 +0000688"""
Barry Warsaw04f357c2002-07-23 19:04:11 +0000689# Emacs turd '
Tim Peters6ba5f792001-06-23 20:45:43 +0000690
Tim Peters0f9da0a2001-06-23 21:01:47 +0000691# Fun tests (for sufficiently warped notions of "fun").
692
693fun_tests = """
694
695Build up to a recursive Sieve of Eratosthenes generator.
696
697>>> def firstn(g, n):
Georg Brandla18af4e2007-04-21 15:47:16 +0000698... return [next(g) for i in range(n)]
Tim Peters0f9da0a2001-06-23 21:01:47 +0000699
700>>> def intsfrom(i):
701... while 1:
702... yield i
703... i += 1
704
705>>> firstn(intsfrom(5), 7)
706[5, 6, 7, 8, 9, 10, 11]
707
708>>> def exclude_multiples(n, ints):
709... for i in ints:
710... if i % n:
711... yield i
712
713>>> firstn(exclude_multiples(3, intsfrom(1)), 6)
714[1, 2, 4, 5, 7, 8]
715
716>>> def sieve(ints):
Georg Brandla18af4e2007-04-21 15:47:16 +0000717... prime = next(ints)
Tim Peters0f9da0a2001-06-23 21:01:47 +0000718... yield prime
719... not_divisible_by_prime = exclude_multiples(prime, ints)
720... for p in sieve(not_divisible_by_prime):
721... yield p
722
723>>> primes = sieve(intsfrom(2))
724>>> firstn(primes, 20)
725[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 +0000726
Tim Petersf6ed0742001-06-27 07:17:57 +0000727
Tim Petersb9e9ff12001-06-24 03:44:52 +0000728Another famous problem: generate all integers of the form
729 2**i * 3**j * 5**k
730in increasing order, where i,j,k >= 0. Trickier than it may look at first!
731Try writing it without generators, and correctly, and without generating
7323 internal results for each result output.
733
734>>> def times(n, g):
735... for i in g:
736... yield n * i
737>>> firstn(times(10, intsfrom(1)), 10)
738[10, 20, 30, 40, 50, 60, 70, 80, 90, 100]
739
740>>> def merge(g, h):
Georg Brandla18af4e2007-04-21 15:47:16 +0000741... ng = next(g)
742... nh = next(h)
Tim Petersb9e9ff12001-06-24 03:44:52 +0000743... while 1:
744... if ng < nh:
745... yield ng
Georg Brandla18af4e2007-04-21 15:47:16 +0000746... ng = next(g)
Tim Petersb9e9ff12001-06-24 03:44:52 +0000747... elif ng > nh:
748... yield nh
Georg Brandla18af4e2007-04-21 15:47:16 +0000749... nh = next(h)
Tim Petersb9e9ff12001-06-24 03:44:52 +0000750... else:
751... yield ng
Georg Brandla18af4e2007-04-21 15:47:16 +0000752... ng = next(g)
753... nh = next(h)
Tim Petersb9e9ff12001-06-24 03:44:52 +0000754
Tim Petersf6ed0742001-06-27 07:17:57 +0000755The following works, but is doing a whale of a lot of redundant work --
756it's not clear how to get the internal uses of m235 to share a single
757generator. Note that me_times2 (etc) each need to see every element in the
758result sequence. So this is an example where lazy lists are more natural
759(you can look at the head of a lazy list any number of times).
Tim Petersb9e9ff12001-06-24 03:44:52 +0000760
761>>> def m235():
762... yield 1
763... me_times2 = times(2, m235())
764... me_times3 = times(3, m235())
765... me_times5 = times(5, m235())
766... for i in merge(merge(me_times2,
767... me_times3),
768... me_times5):
769... yield i
770
Tim Petersf6ed0742001-06-27 07:17:57 +0000771Don't print "too many" of these -- the implementation above is extremely
772inefficient: each call of m235() leads to 3 recursive calls, and in
773turn each of those 3 more, and so on, and so on, until we've descended
774enough levels to satisfy the print stmts. Very odd: when I printed 5
775lines of results below, this managed to screw up Win98's malloc in "the
776usual" way, i.e. the heap grew over 4Mb so Win98 started fragmenting
777address space, and it *looked* like a very slow leak.
778
Tim Petersb9e9ff12001-06-24 03:44:52 +0000779>>> result = m235()
Tim Petersf6ed0742001-06-27 07:17:57 +0000780>>> for i in range(3):
Guido van Rossum7131f842007-02-09 20:13:25 +0000781... print(firstn(result, 15))
Tim Petersb9e9ff12001-06-24 03:44:52 +0000782[1, 2, 3, 4, 5, 6, 8, 9, 10, 12, 15, 16, 18, 20, 24]
783[25, 27, 30, 32, 36, 40, 45, 48, 50, 54, 60, 64, 72, 75, 80]
784[81, 90, 96, 100, 108, 120, 125, 128, 135, 144, 150, 160, 162, 180, 192]
Tim Petersee309272001-06-24 05:47:06 +0000785
786Heh. Here's one way to get a shared list, complete with an excruciating
787namespace renaming trick. The *pretty* part is that the times() and merge()
788functions can be reused as-is, because they only assume their stream
789arguments are iterable -- a LazyList is the same as a generator to times().
790
791>>> class LazyList:
792... def __init__(self, g):
793... self.sofar = []
Georg Brandla18af4e2007-04-21 15:47:16 +0000794... self.fetch = g.__next__
Tim Petersee309272001-06-24 05:47:06 +0000795...
796... def __getitem__(self, i):
797... sofar, fetch = self.sofar, self.fetch
798... while i >= len(sofar):
799... sofar.append(fetch())
800... return sofar[i]
801
802>>> def m235():
803... yield 1
Tim Petersea2e97a2001-06-24 07:10:02 +0000804... # Gack: m235 below actually refers to a LazyList.
Tim Petersee309272001-06-24 05:47:06 +0000805... me_times2 = times(2, m235)
806... me_times3 = times(3, m235)
807... me_times5 = times(5, m235)
808... for i in merge(merge(me_times2,
809... me_times3),
810... me_times5):
811... yield i
812
Tim Petersf6ed0742001-06-27 07:17:57 +0000813Print as many of these as you like -- *this* implementation is memory-
Neil Schemenauerb20e9db2001-07-12 13:26:41 +0000814efficient.
Tim Petersf6ed0742001-06-27 07:17:57 +0000815
Tim Petersee309272001-06-24 05:47:06 +0000816>>> m235 = LazyList(m235())
817>>> for i in range(5):
Guido van Rossum7131f842007-02-09 20:13:25 +0000818... print([m235[j] for j in range(15*i, 15*(i+1))])
Tim Petersee309272001-06-24 05:47:06 +0000819[1, 2, 3, 4, 5, 6, 8, 9, 10, 12, 15, 16, 18, 20, 24]
820[25, 27, 30, 32, 36, 40, 45, 48, 50, 54, 60, 64, 72, 75, 80]
821[81, 90, 96, 100, 108, 120, 125, 128, 135, 144, 150, 160, 162, 180, 192]
822[200, 216, 225, 240, 243, 250, 256, 270, 288, 300, 320, 324, 360, 375, 384]
823[400, 405, 432, 450, 480, 486, 500, 512, 540, 576, 600, 625, 640, 648, 675]
Tim Petersf6ed0742001-06-27 07:17:57 +0000824
Tim Petersf6ed0742001-06-27 07:17:57 +0000825Ye olde Fibonacci generator, LazyList style.
826
827>>> def fibgen(a, b):
828...
829... def sum(g, h):
830... while 1:
Georg Brandla18af4e2007-04-21 15:47:16 +0000831... yield next(g) + next(h)
Tim Petersf6ed0742001-06-27 07:17:57 +0000832...
833... def tail(g):
Georg Brandla18af4e2007-04-21 15:47:16 +0000834... next(g) # throw first away
Tim Petersf6ed0742001-06-27 07:17:57 +0000835... for x in g:
836... yield x
837...
838... yield a
839... yield b
840... for s in sum(iter(fib),
841... tail(iter(fib))):
842... yield s
843
844>>> fib = LazyList(fibgen(1, 2))
845>>> firstn(iter(fib), 17)
846[1, 2, 3, 5, 8, 13, 21, 34, 55, 89, 144, 233, 377, 610, 987, 1597, 2584]
Georg Brandl52715f62005-08-24 09:02:29 +0000847
848
849Running after your tail with itertools.tee (new in version 2.4)
850
851The algorithms "m235" (Hamming) and Fibonacci presented above are both
852examples of a whole family of FP (functional programming) algorithms
853where a function produces and returns a list while the production algorithm
854suppose the list as already produced by recursively calling itself.
855For these algorithms to work, they must:
856
857- produce at least a first element without presupposing the existence of
858 the rest of the list
859- produce their elements in a lazy manner
860
861To work efficiently, the beginning of the list must not be recomputed over
862and over again. This is ensured in most FP languages as a built-in feature.
863In python, we have to explicitly maintain a list of already computed results
864and abandon genuine recursivity.
865
866This is what had been attempted above with the LazyList class. One problem
867with that class is that it keeps a list of all of the generated results and
868therefore continually grows. This partially defeats the goal of the generator
869concept, viz. produce the results only as needed instead of producing them
870all and thereby wasting memory.
871
872Thanks to itertools.tee, it is now clear "how to get the internal uses of
873m235 to share a single generator".
874
875>>> from itertools import tee
876>>> def m235():
877... def _m235():
878... yield 1
879... for n in merge(times(2, m2),
880... merge(times(3, m3),
881... times(5, m5))):
882... yield n
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000883... m1 = _m235()
884... m2, m3, m5, mRes = tee(m1, 4)
Georg Brandl52715f62005-08-24 09:02:29 +0000885... return mRes
886
887>>> it = m235()
888>>> for i in range(5):
Guido van Rossum7131f842007-02-09 20:13:25 +0000889... print(firstn(it, 15))
Georg Brandl52715f62005-08-24 09:02:29 +0000890[1, 2, 3, 4, 5, 6, 8, 9, 10, 12, 15, 16, 18, 20, 24]
891[25, 27, 30, 32, 36, 40, 45, 48, 50, 54, 60, 64, 72, 75, 80]
892[81, 90, 96, 100, 108, 120, 125, 128, 135, 144, 150, 160, 162, 180, 192]
893[200, 216, 225, 240, 243, 250, 256, 270, 288, 300, 320, 324, 360, 375, 384]
894[400, 405, 432, 450, 480, 486, 500, 512, 540, 576, 600, 625, 640, 648, 675]
895
896The "tee" function does just what we want. It internally keeps a generated
897result for as long as it has not been "consumed" from all of the duplicated
898iterators, whereupon it is deleted. You can therefore print the hamming
899sequence during hours without increasing memory usage, or very little.
900
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000901The beauty of it is that recursive running-after-their-tail FP algorithms
Georg Brandl52715f62005-08-24 09:02:29 +0000902are quite straightforwardly expressed with this Python idiom.
903
Georg Brandl52715f62005-08-24 09:02:29 +0000904Ye olde Fibonacci generator, tee style.
905
906>>> def fib():
Tim Peters9e34c042005-08-26 15:20:46 +0000907...
Georg Brandl52715f62005-08-24 09:02:29 +0000908... def _isum(g, h):
909... while 1:
Georg Brandla18af4e2007-04-21 15:47:16 +0000910... yield next(g) + next(h)
Tim Peters9e34c042005-08-26 15:20:46 +0000911...
Georg Brandl52715f62005-08-24 09:02:29 +0000912... def _fib():
913... yield 1
914... yield 2
Georg Brandla18af4e2007-04-21 15:47:16 +0000915... next(fibTail) # throw first away
Georg Brandl52715f62005-08-24 09:02:29 +0000916... for res in _isum(fibHead, fibTail):
917... yield res
Tim Peters9e34c042005-08-26 15:20:46 +0000918...
Thomas Wouters49fd7fa2006-04-21 10:40:58 +0000919... realfib = _fib()
920... fibHead, fibTail, fibRes = tee(realfib, 3)
Georg Brandl52715f62005-08-24 09:02:29 +0000921... return fibRes
922
923>>> firstn(fib(), 17)
924[1, 2, 3, 5, 8, 13, 21, 34, 55, 89, 144, 233, 377, 610, 987, 1597, 2584]
925
Tim Peters0f9da0a2001-06-23 21:01:47 +0000926"""
927
Tim Petersb6c3cea2001-06-26 03:36:28 +0000928# syntax_tests mostly provokes SyntaxErrors. Also fiddling with #if 0
929# hackery.
Tim Petersee309272001-06-24 05:47:06 +0000930
Tim Petersea2e97a2001-06-24 07:10:02 +0000931syntax_tests = """
932
Phillip J. Eby0d6615f2005-08-02 00:46:46 +0000933These are fine:
Tim Petersea2e97a2001-06-24 07:10:02 +0000934
935>>> def f():
936... yield 1
937... return
938
Tim Petersaef8cfa2004-08-27 15:12:49 +0000939>>> def f():
Tim Petersea2e97a2001-06-24 07:10:02 +0000940... try:
941... yield 1
942... finally:
943... pass
Tim Petersea2e97a2001-06-24 07:10:02 +0000944
Tim Petersaef8cfa2004-08-27 15:12:49 +0000945>>> def f():
Tim Petersea2e97a2001-06-24 07:10:02 +0000946... try:
947... try:
Tim Peters3caca232001-12-06 06:23:26 +0000948... 1//0
Tim Petersea2e97a2001-06-24 07:10:02 +0000949... except ZeroDivisionError:
Tim Peters536cf992005-12-25 23:18:31 +0000950... yield 666
Tim Petersea2e97a2001-06-24 07:10:02 +0000951... except:
952... pass
953... finally:
954... pass
Tim Petersea2e97a2001-06-24 07:10:02 +0000955
956>>> def f():
957... try:
958... try:
959... yield 12
Tim Peters3caca232001-12-06 06:23:26 +0000960... 1//0
Tim Petersea2e97a2001-06-24 07:10:02 +0000961... except ZeroDivisionError:
962... yield 666
963... except:
964... try:
965... x = 12
966... finally:
967... yield 12
968... except:
969... return
970>>> list(f())
971[12, 666]
Tim Petersb6c3cea2001-06-26 03:36:28 +0000972
973>>> def f():
Tim Peters08a898f2001-06-28 01:52:22 +0000974... yield
Phillip J. Eby0d6615f2005-08-02 00:46:46 +0000975>>> type(f())
Martin v. Löwis250ad612008-04-07 05:43:42 +0000976<class 'generator'>
Phillip J. Eby0d6615f2005-08-02 00:46:46 +0000977
Tim Peters08a898f2001-06-28 01:52:22 +0000978
979>>> def f():
980... if 0:
981... yield
Phillip J. Eby0d6615f2005-08-02 00:46:46 +0000982>>> type(f())
Martin v. Löwis250ad612008-04-07 05:43:42 +0000983<class 'generator'>
Phillip J. Eby0d6615f2005-08-02 00:46:46 +0000984
Tim Peters08a898f2001-06-28 01:52:22 +0000985
986>>> def f():
Tim Petersb6c3cea2001-06-26 03:36:28 +0000987... if 0:
988... yield 1
989>>> type(f())
Martin v. Löwis250ad612008-04-07 05:43:42 +0000990<class 'generator'>
Tim Petersb6c3cea2001-06-26 03:36:28 +0000991
992>>> def f():
993... if "":
994... yield None
995>>> type(f())
Martin v. Löwis250ad612008-04-07 05:43:42 +0000996<class 'generator'>
Tim Petersb6c3cea2001-06-26 03:36:28 +0000997
998>>> def f():
999... return
1000... try:
1001... if x==4:
1002... pass
1003... elif 0:
1004... try:
Tim Peters3caca232001-12-06 06:23:26 +00001005... 1//0
Tim Petersb6c3cea2001-06-26 03:36:28 +00001006... except SyntaxError:
1007... pass
1008... else:
1009... if 0:
1010... while 12:
1011... x += 1
1012... yield 2 # don't blink
1013... f(a, b, c, d, e)
1014... else:
1015... pass
1016... except:
1017... x = 1
1018... return
1019>>> type(f())
Martin v. Löwis250ad612008-04-07 05:43:42 +00001020<class 'generator'>
Tim Petersb6c3cea2001-06-26 03:36:28 +00001021
1022>>> def f():
1023... if 0:
1024... def g():
1025... yield 1
1026...
1027>>> type(f())
Martin v. Löwis250ad612008-04-07 05:43:42 +00001028<class 'NoneType'>
Tim Petersb6c3cea2001-06-26 03:36:28 +00001029
1030>>> def f():
1031... if 0:
1032... class C:
1033... def __init__(self):
1034... yield 1
1035... def f(self):
1036... yield 2
1037>>> type(f())
Martin v. Löwis250ad612008-04-07 05:43:42 +00001038<class 'NoneType'>
Tim Peters08a898f2001-06-28 01:52:22 +00001039
1040>>> def f():
1041... if 0:
1042... return
1043... if 0:
1044... yield 2
1045>>> type(f())
Martin v. Löwis250ad612008-04-07 05:43:42 +00001046<class 'generator'>
Tim Peters08a898f2001-06-28 01:52:22 +00001047
Guido van Rossumc5fe5eb2002-06-12 03:45:21 +00001048This one caused a crash (see SF bug 567538):
1049
1050>>> def f():
1051... for i in range(3):
1052... try:
1053... continue
1054... finally:
1055... yield i
Tim Petersc411dba2002-07-16 21:35:23 +00001056...
Guido van Rossumc5fe5eb2002-06-12 03:45:21 +00001057>>> g = f()
Georg Brandla18af4e2007-04-21 15:47:16 +00001058>>> print(next(g))
Guido van Rossumc5fe5eb2002-06-12 03:45:21 +000010590
Georg Brandla18af4e2007-04-21 15:47:16 +00001060>>> print(next(g))
Guido van Rossumc5fe5eb2002-06-12 03:45:21 +000010611
Georg Brandla18af4e2007-04-21 15:47:16 +00001062>>> print(next(g))
Guido van Rossumc5fe5eb2002-06-12 03:45:21 +000010632
Georg Brandla18af4e2007-04-21 15:47:16 +00001064>>> print(next(g))
Guido van Rossumc5fe5eb2002-06-12 03:45:21 +00001065Traceback (most recent call last):
1066StopIteration
Christian Heimesaf98da12008-01-27 15:18:18 +00001067
1068
1069Test the gi_code attribute
1070
1071>>> def f():
1072... yield 5
1073...
1074>>> g = f()
1075>>> g.gi_code is f.__code__
1076True
1077>>> next(g)
10785
1079>>> next(g)
1080Traceback (most recent call last):
1081StopIteration
1082>>> g.gi_code is f.__code__
1083True
1084
Alexandre Vassalottie9f305f2008-05-16 04:39:54 +00001085
1086Test the __name__ attribute and the repr()
1087
1088>>> def f():
1089... yield 5
1090...
1091>>> g = f()
1092>>> g.__name__
1093'f'
1094>>> repr(g) # doctest: +ELLIPSIS
Alexandre Vassalottibee32532008-05-16 18:15:12 +00001095'<generator object f at ...>'
Benjamin Peterson371ccfb2008-12-27 19:03:36 +00001096
1097Lambdas shouldn't have their usual return behavior.
1098
1099>>> x = lambda: (yield 1)
1100>>> list(x())
1101[1]
1102
1103>>> x = lambda: ((yield 1), (yield 2))
1104>>> list(x())
1105[1, 2]
Tim Petersea2e97a2001-06-24 07:10:02 +00001106"""
1107
Tim Petersbe4f0a72001-06-29 02:41:16 +00001108# conjoin is a simple backtracking generator, named in honor of Icon's
1109# "conjunction" control structure. Pass a list of no-argument functions
1110# that return iterable objects. Easiest to explain by example: assume the
1111# function list [x, y, z] is passed. Then conjoin acts like:
1112#
1113# def g():
1114# values = [None] * 3
1115# for values[0] in x():
1116# for values[1] in y():
1117# for values[2] in z():
1118# yield values
1119#
1120# So some 3-lists of values *may* be generated, each time we successfully
1121# get into the innermost loop. If an iterator fails (is exhausted) before
1122# then, it "backtracks" to get the next value from the nearest enclosing
1123# iterator (the one "to the left"), and starts all over again at the next
1124# slot (pumps a fresh iterator). Of course this is most useful when the
1125# iterators have side-effects, so that which values *can* be generated at
1126# each slot depend on the values iterated at previous slots.
1127
Benjamin Peterson78565b22009-06-28 19:19:51 +00001128def simple_conjoin(gs):
Tim Petersbe4f0a72001-06-29 02:41:16 +00001129
1130 values = [None] * len(gs)
1131
Benjamin Peterson78565b22009-06-28 19:19:51 +00001132 def gen(i):
Tim Petersbe4f0a72001-06-29 02:41:16 +00001133 if i >= len(gs):
1134 yield values
1135 else:
1136 for values[i] in gs[i]():
1137 for x in gen(i+1):
1138 yield x
1139
1140 for x in gen(0):
1141 yield x
1142
Tim Petersc468fd22001-06-30 07:29:44 +00001143# That works fine, but recursing a level and checking i against len(gs) for
1144# each item produced is inefficient. By doing manual loop unrolling across
1145# generator boundaries, it's possible to eliminate most of that overhead.
1146# This isn't worth the bother *in general* for generators, but conjoin() is
1147# a core building block for some CPU-intensive generator applications.
1148
1149def conjoin(gs):
1150
1151 n = len(gs)
1152 values = [None] * n
1153
1154 # Do one loop nest at time recursively, until the # of loop nests
1155 # remaining is divisible by 3.
1156
Benjamin Peterson78565b22009-06-28 19:19:51 +00001157 def gen(i):
Tim Petersc468fd22001-06-30 07:29:44 +00001158 if i >= n:
1159 yield values
1160
1161 elif (n-i) % 3:
1162 ip1 = i+1
1163 for values[i] in gs[i]():
1164 for x in gen(ip1):
1165 yield x
1166
1167 else:
1168 for x in _gen3(i):
1169 yield x
1170
1171 # Do three loop nests at a time, recursing only if at least three more
1172 # remain. Don't call directly: this is an internal optimization for
1173 # gen's use.
1174
Benjamin Peterson78565b22009-06-28 19:19:51 +00001175 def _gen3(i):
Tim Petersc468fd22001-06-30 07:29:44 +00001176 assert i < n and (n-i) % 3 == 0
1177 ip1, ip2, ip3 = i+1, i+2, i+3
1178 g, g1, g2 = gs[i : ip3]
1179
1180 if ip3 >= n:
1181 # These are the last three, so we can yield values directly.
1182 for values[i] in g():
1183 for values[ip1] in g1():
1184 for values[ip2] in g2():
1185 yield values
1186
1187 else:
1188 # At least 6 loop nests remain; peel off 3 and recurse for the
1189 # rest.
1190 for values[i] in g():
1191 for values[ip1] in g1():
1192 for values[ip2] in g2():
1193 for x in _gen3(ip3):
1194 yield x
1195
1196 for x in gen(0):
1197 yield x
1198
unknown31569562001-07-04 22:11:22 +00001199# And one more approach: For backtracking apps like the Knight's Tour
1200# solver below, the number of backtracking levels can be enormous (one
1201# level per square, for the Knight's Tour, so that e.g. a 100x100 board
1202# needs 10,000 levels). In such cases Python is likely to run out of
1203# stack space due to recursion. So here's a recursion-free version of
1204# conjoin too.
1205# NOTE WELL: This allows large problems to be solved with only trivial
1206# demands on stack space. Without explicitly resumable generators, this is
Tim Peters9a8c8e22001-07-13 09:12:12 +00001207# much harder to achieve. OTOH, this is much slower (up to a factor of 2)
1208# than the fancy unrolled recursive conjoin.
unknown31569562001-07-04 22:11:22 +00001209
1210def flat_conjoin(gs): # rename to conjoin to run tests with this instead
1211 n = len(gs)
1212 values = [None] * n
1213 iters = [None] * n
Tim Peters9a8c8e22001-07-13 09:12:12 +00001214 _StopIteration = StopIteration # make local because caught a *lot*
unknown31569562001-07-04 22:11:22 +00001215 i = 0
Tim Peters9a8c8e22001-07-13 09:12:12 +00001216 while 1:
1217 # Descend.
1218 try:
1219 while i < n:
Georg Brandla18af4e2007-04-21 15:47:16 +00001220 it = iters[i] = gs[i]().__next__
Tim Peters9a8c8e22001-07-13 09:12:12 +00001221 values[i] = it()
1222 i += 1
1223 except _StopIteration:
1224 pass
unknown31569562001-07-04 22:11:22 +00001225 else:
Tim Peters9a8c8e22001-07-13 09:12:12 +00001226 assert i == n
1227 yield values
unknown31569562001-07-04 22:11:22 +00001228
Tim Peters9a8c8e22001-07-13 09:12:12 +00001229 # Backtrack until an older iterator can be resumed.
1230 i -= 1
unknown31569562001-07-04 22:11:22 +00001231 while i >= 0:
1232 try:
1233 values[i] = iters[i]()
Tim Peters9a8c8e22001-07-13 09:12:12 +00001234 # Success! Start fresh at next level.
unknown31569562001-07-04 22:11:22 +00001235 i += 1
1236 break
Tim Peters9a8c8e22001-07-13 09:12:12 +00001237 except _StopIteration:
1238 # Continue backtracking.
1239 i -= 1
1240 else:
1241 assert i < 0
1242 break
unknown31569562001-07-04 22:11:22 +00001243
Tim Petersbe4f0a72001-06-29 02:41:16 +00001244# A conjoin-based N-Queens solver.
1245
1246class Queens:
1247 def __init__(self, n):
1248 self.n = n
1249 rangen = range(n)
1250
1251 # Assign a unique int to each column and diagonal.
1252 # columns: n of those, range(n).
1253 # NW-SE diagonals: 2n-1 of these, i-j unique and invariant along
1254 # each, smallest i-j is 0-(n-1) = 1-n, so add n-1 to shift to 0-
1255 # based.
1256 # NE-SW diagonals: 2n-1 of these, i+j unique and invariant along
1257 # each, smallest i+j is 0, largest is 2n-2.
1258
1259 # For each square, compute a bit vector of the columns and
1260 # diagonals it covers, and for each row compute a function that
1261 # generates the possiblities for the columns in that row.
1262 self.rowgenerators = []
1263 for i in rangen:
Guido van Rossume2a383d2007-01-15 16:59:06 +00001264 rowuses = [(1 << j) | # column ordinal
1265 (1 << (n + i-j + n-1)) | # NW-SE ordinal
1266 (1 << (n + 2*n-1 + i+j)) # NE-SW ordinal
Tim Petersbe4f0a72001-06-29 02:41:16 +00001267 for j in rangen]
1268
1269 def rowgen(rowuses=rowuses):
1270 for j in rangen:
1271 uses = rowuses[j]
Tim Petersc468fd22001-06-30 07:29:44 +00001272 if uses & self.used == 0:
1273 self.used |= uses
1274 yield j
1275 self.used &= ~uses
Tim Petersbe4f0a72001-06-29 02:41:16 +00001276
1277 self.rowgenerators.append(rowgen)
1278
1279 # Generate solutions.
1280 def solve(self):
1281 self.used = 0
1282 for row2col in conjoin(self.rowgenerators):
1283 yield row2col
1284
1285 def printsolution(self, row2col):
1286 n = self.n
1287 assert n == len(row2col)
1288 sep = "+" + "-+" * n
Guido van Rossumbe19ed72007-02-09 05:37:30 +00001289 print(sep)
Tim Petersbe4f0a72001-06-29 02:41:16 +00001290 for i in range(n):
1291 squares = [" " for j in range(n)]
1292 squares[row2col[i]] = "Q"
Guido van Rossumbe19ed72007-02-09 05:37:30 +00001293 print("|" + "|".join(squares) + "|")
1294 print(sep)
Tim Petersbe4f0a72001-06-29 02:41:16 +00001295
unknown31569562001-07-04 22:11:22 +00001296# A conjoin-based Knight's Tour solver. This is pretty sophisticated
1297# (e.g., when used with flat_conjoin above, and passing hard=1 to the
1298# constructor, a 200x200 Knight's Tour was found quickly -- note that we're
Tim Peters9a8c8e22001-07-13 09:12:12 +00001299# creating 10s of thousands of generators then!), and is lengthy.
unknown31569562001-07-04 22:11:22 +00001300
1301class Knights:
Tim Peters9a8c8e22001-07-13 09:12:12 +00001302 def __init__(self, m, n, hard=0):
1303 self.m, self.n = m, n
unknown31569562001-07-04 22:11:22 +00001304
Tim Peters9a8c8e22001-07-13 09:12:12 +00001305 # solve() will set up succs[i] to be a list of square #i's
1306 # successors.
1307 succs = self.succs = []
unknown31569562001-07-04 22:11:22 +00001308
Tim Peters9a8c8e22001-07-13 09:12:12 +00001309 # Remove i0 from each of its successor's successor lists, i.e.
1310 # successors can't go back to i0 again. Return 0 if we can
1311 # detect this makes a solution impossible, else return 1.
unknown31569562001-07-04 22:11:22 +00001312
Tim Peters9a8c8e22001-07-13 09:12:12 +00001313 def remove_from_successors(i0, len=len):
unknown31569562001-07-04 22:11:22 +00001314 # If we remove all exits from a free square, we're dead:
1315 # even if we move to it next, we can't leave it again.
1316 # If we create a square with one exit, we must visit it next;
1317 # else somebody else will have to visit it, and since there's
1318 # only one adjacent, there won't be a way to leave it again.
1319 # Finelly, if we create more than one free square with a
1320 # single exit, we can only move to one of them next, leaving
1321 # the other one a dead end.
1322 ne0 = ne1 = 0
1323 for i in succs[i0]:
Tim Peters9a8c8e22001-07-13 09:12:12 +00001324 s = succs[i]
1325 s.remove(i0)
1326 e = len(s)
1327 if e == 0:
1328 ne0 += 1
1329 elif e == 1:
1330 ne1 += 1
unknown31569562001-07-04 22:11:22 +00001331 return ne0 == 0 and ne1 < 2
1332
Tim Peters9a8c8e22001-07-13 09:12:12 +00001333 # Put i0 back in each of its successor's successor lists.
1334
1335 def add_to_successors(i0):
unknown31569562001-07-04 22:11:22 +00001336 for i in succs[i0]:
Tim Peters9a8c8e22001-07-13 09:12:12 +00001337 succs[i].append(i0)
unknown31569562001-07-04 22:11:22 +00001338
1339 # Generate the first move.
1340 def first():
Tim Peters9a8c8e22001-07-13 09:12:12 +00001341 if m < 1 or n < 1:
unknown31569562001-07-04 22:11:22 +00001342 return
1343
unknown31569562001-07-04 22:11:22 +00001344 # Since we're looking for a cycle, it doesn't matter where we
1345 # start. Starting in a corner makes the 2nd move easy.
Tim Peters9a8c8e22001-07-13 09:12:12 +00001346 corner = self.coords2index(0, 0)
1347 remove_from_successors(corner)
unknown31569562001-07-04 22:11:22 +00001348 self.lastij = corner
1349 yield corner
Tim Peters9a8c8e22001-07-13 09:12:12 +00001350 add_to_successors(corner)
unknown31569562001-07-04 22:11:22 +00001351
1352 # Generate the second moves.
1353 def second():
Tim Peters9a8c8e22001-07-13 09:12:12 +00001354 corner = self.coords2index(0, 0)
unknown31569562001-07-04 22:11:22 +00001355 assert self.lastij == corner # i.e., we started in the corner
Tim Peters9a8c8e22001-07-13 09:12:12 +00001356 if m < 3 or n < 3:
unknown31569562001-07-04 22:11:22 +00001357 return
Tim Peters9a8c8e22001-07-13 09:12:12 +00001358 assert len(succs[corner]) == 2
1359 assert self.coords2index(1, 2) in succs[corner]
1360 assert self.coords2index(2, 1) in succs[corner]
unknown31569562001-07-04 22:11:22 +00001361 # Only two choices. Whichever we pick, the other must be the
Tim Peters9a8c8e22001-07-13 09:12:12 +00001362 # square picked on move m*n, as it's the only way to get back
unknown31569562001-07-04 22:11:22 +00001363 # to (0, 0). Save its index in self.final so that moves before
1364 # the last know it must be kept free.
1365 for i, j in (1, 2), (2, 1):
Tim Peters9a8c8e22001-07-13 09:12:12 +00001366 this = self.coords2index(i, j)
1367 final = self.coords2index(3-i, 3-j)
unknown31569562001-07-04 22:11:22 +00001368 self.final = final
unknown31569562001-07-04 22:11:22 +00001369
Tim Peters9a8c8e22001-07-13 09:12:12 +00001370 remove_from_successors(this)
1371 succs[final].append(corner)
unknown31569562001-07-04 22:11:22 +00001372 self.lastij = this
1373 yield this
Tim Peters9a8c8e22001-07-13 09:12:12 +00001374 succs[final].remove(corner)
1375 add_to_successors(this)
unknown31569562001-07-04 22:11:22 +00001376
Tim Peters9a8c8e22001-07-13 09:12:12 +00001377 # Generate moves 3 thru m*n-1.
1378 def advance(len=len):
unknown31569562001-07-04 22:11:22 +00001379 # If some successor has only one exit, must take it.
1380 # Else favor successors with fewer exits.
1381 candidates = []
1382 for i in succs[self.lastij]:
Tim Peters9a8c8e22001-07-13 09:12:12 +00001383 e = len(succs[i])
1384 assert e > 0, "else remove_from_successors() pruning flawed"
1385 if e == 1:
1386 candidates = [(e, i)]
1387 break
1388 candidates.append((e, i))
unknown31569562001-07-04 22:11:22 +00001389 else:
1390 candidates.sort()
1391
1392 for e, i in candidates:
1393 if i != self.final:
Tim Peters9a8c8e22001-07-13 09:12:12 +00001394 if remove_from_successors(i):
unknown31569562001-07-04 22:11:22 +00001395 self.lastij = i
1396 yield i
Tim Peters9a8c8e22001-07-13 09:12:12 +00001397 add_to_successors(i)
unknown31569562001-07-04 22:11:22 +00001398
Tim Peters9a8c8e22001-07-13 09:12:12 +00001399 # Generate moves 3 thru m*n-1. Alternative version using a
unknown31569562001-07-04 22:11:22 +00001400 # stronger (but more expensive) heuristic to order successors.
Tim Peters9a8c8e22001-07-13 09:12:12 +00001401 # Since the # of backtracking levels is m*n, a poor move early on
1402 # can take eons to undo. Smallest square board for which this
1403 # matters a lot is 52x52.
1404 def advance_hard(vmid=(m-1)/2.0, hmid=(n-1)/2.0, len=len):
unknown31569562001-07-04 22:11:22 +00001405 # If some successor has only one exit, must take it.
1406 # Else favor successors with fewer exits.
1407 # Break ties via max distance from board centerpoint (favor
1408 # corners and edges whenever possible).
1409 candidates = []
1410 for i in succs[self.lastij]:
Tim Peters9a8c8e22001-07-13 09:12:12 +00001411 e = len(succs[i])
1412 assert e > 0, "else remove_from_successors() pruning flawed"
1413 if e == 1:
1414 candidates = [(e, 0, i)]
1415 break
1416 i1, j1 = self.index2coords(i)
1417 d = (i1 - vmid)**2 + (j1 - hmid)**2
1418 candidates.append((e, -d, i))
unknown31569562001-07-04 22:11:22 +00001419 else:
1420 candidates.sort()
1421
1422 for e, d, i in candidates:
1423 if i != self.final:
Tim Peters9a8c8e22001-07-13 09:12:12 +00001424 if remove_from_successors(i):
unknown31569562001-07-04 22:11:22 +00001425 self.lastij = i
1426 yield i
Tim Peters9a8c8e22001-07-13 09:12:12 +00001427 add_to_successors(i)
unknown31569562001-07-04 22:11:22 +00001428
1429 # Generate the last move.
1430 def last():
1431 assert self.final in succs[self.lastij]
1432 yield self.final
1433
Tim Peters9a8c8e22001-07-13 09:12:12 +00001434 if m*n < 4:
1435 self.squaregenerators = [first]
unknown31569562001-07-04 22:11:22 +00001436 else:
Tim Peters9a8c8e22001-07-13 09:12:12 +00001437 self.squaregenerators = [first, second] + \
1438 [hard and advance_hard or advance] * (m*n - 3) + \
unknown31569562001-07-04 22:11:22 +00001439 [last]
1440
Tim Peters9a8c8e22001-07-13 09:12:12 +00001441 def coords2index(self, i, j):
1442 assert 0 <= i < self.m
1443 assert 0 <= j < self.n
1444 return i * self.n + j
1445
1446 def index2coords(self, index):
1447 assert 0 <= index < self.m * self.n
1448 return divmod(index, self.n)
1449
1450 def _init_board(self):
1451 succs = self.succs
1452 del succs[:]
1453 m, n = self.m, self.n
1454 c2i = self.coords2index
1455
1456 offsets = [( 1, 2), ( 2, 1), ( 2, -1), ( 1, -2),
1457 (-1, -2), (-2, -1), (-2, 1), (-1, 2)]
1458 rangen = range(n)
1459 for i in range(m):
1460 for j in rangen:
1461 s = [c2i(i+io, j+jo) for io, jo in offsets
1462 if 0 <= i+io < m and
1463 0 <= j+jo < n]
1464 succs.append(s)
1465
unknown31569562001-07-04 22:11:22 +00001466 # Generate solutions.
1467 def solve(self):
Tim Peters9a8c8e22001-07-13 09:12:12 +00001468 self._init_board()
1469 for x in conjoin(self.squaregenerators):
unknown31569562001-07-04 22:11:22 +00001470 yield x
1471
1472 def printsolution(self, x):
Tim Peters9a8c8e22001-07-13 09:12:12 +00001473 m, n = self.m, self.n
1474 assert len(x) == m*n
1475 w = len(str(m*n))
unknown31569562001-07-04 22:11:22 +00001476 format = "%" + str(w) + "d"
1477
Tim Peters9a8c8e22001-07-13 09:12:12 +00001478 squares = [[None] * n for i in range(m)]
unknown31569562001-07-04 22:11:22 +00001479 k = 1
1480 for i in x:
Tim Peters9a8c8e22001-07-13 09:12:12 +00001481 i1, j1 = self.index2coords(i)
unknown31569562001-07-04 22:11:22 +00001482 squares[i1][j1] = format % k
1483 k += 1
1484
1485 sep = "+" + ("-" * w + "+") * n
Guido van Rossumbe19ed72007-02-09 05:37:30 +00001486 print(sep)
Tim Peters9a8c8e22001-07-13 09:12:12 +00001487 for i in range(m):
unknown31569562001-07-04 22:11:22 +00001488 row = squares[i]
Guido van Rossumbe19ed72007-02-09 05:37:30 +00001489 print("|" + "|".join(row) + "|")
1490 print(sep)
unknown31569562001-07-04 22:11:22 +00001491
Tim Petersbe4f0a72001-06-29 02:41:16 +00001492conjoin_tests = """
1493
1494Generate the 3-bit binary numbers in order. This illustrates dumbest-
1495possible use of conjoin, just to generate the full cross-product.
1496
unknown31569562001-07-04 22:11:22 +00001497>>> for c in conjoin([lambda: iter((0, 1))] * 3):
Guido van Rossum7131f842007-02-09 20:13:25 +00001498... print(c)
Tim Petersbe4f0a72001-06-29 02:41:16 +00001499[0, 0, 0]
1500[0, 0, 1]
1501[0, 1, 0]
1502[0, 1, 1]
1503[1, 0, 0]
1504[1, 0, 1]
1505[1, 1, 0]
1506[1, 1, 1]
1507
Tim Petersc468fd22001-06-30 07:29:44 +00001508For efficiency in typical backtracking apps, conjoin() yields the same list
1509object each time. So if you want to save away a full account of its
1510generated sequence, you need to copy its results.
1511
1512>>> def gencopy(iterator):
1513... for x in iterator:
1514... yield x[:]
1515
1516>>> for n in range(10):
unknown31569562001-07-04 22:11:22 +00001517... all = list(gencopy(conjoin([lambda: iter((0, 1))] * n)))
Guido van Rossum7131f842007-02-09 20:13:25 +00001518... print(n, len(all), all[0] == [0] * n, all[-1] == [1] * n)
Guido van Rossum77f6a652002-04-03 22:41:51 +000015190 1 True True
15201 2 True True
15212 4 True True
15223 8 True True
15234 16 True True
15245 32 True True
15256 64 True True
15267 128 True True
15278 256 True True
15289 512 True True
Tim Petersc468fd22001-06-30 07:29:44 +00001529
Tim Petersbe4f0a72001-06-29 02:41:16 +00001530And run an 8-queens solver.
1531
1532>>> q = Queens(8)
1533>>> LIMIT = 2
1534>>> count = 0
1535>>> for row2col in q.solve():
1536... count += 1
1537... if count <= LIMIT:
Guido van Rossum7131f842007-02-09 20:13:25 +00001538... print("Solution", count)
Tim Petersbe4f0a72001-06-29 02:41:16 +00001539... q.printsolution(row2col)
1540Solution 1
1541+-+-+-+-+-+-+-+-+
1542|Q| | | | | | | |
1543+-+-+-+-+-+-+-+-+
1544| | | | |Q| | | |
1545+-+-+-+-+-+-+-+-+
1546| | | | | | | |Q|
1547+-+-+-+-+-+-+-+-+
1548| | | | | |Q| | |
1549+-+-+-+-+-+-+-+-+
1550| | |Q| | | | | |
1551+-+-+-+-+-+-+-+-+
1552| | | | | | |Q| |
1553+-+-+-+-+-+-+-+-+
1554| |Q| | | | | | |
1555+-+-+-+-+-+-+-+-+
1556| | | |Q| | | | |
1557+-+-+-+-+-+-+-+-+
1558Solution 2
1559+-+-+-+-+-+-+-+-+
1560|Q| | | | | | | |
1561+-+-+-+-+-+-+-+-+
1562| | | | | |Q| | |
1563+-+-+-+-+-+-+-+-+
1564| | | | | | | |Q|
1565+-+-+-+-+-+-+-+-+
1566| | |Q| | | | | |
1567+-+-+-+-+-+-+-+-+
1568| | | | | | |Q| |
1569+-+-+-+-+-+-+-+-+
1570| | | |Q| | | | |
1571+-+-+-+-+-+-+-+-+
1572| |Q| | | | | | |
1573+-+-+-+-+-+-+-+-+
1574| | | | |Q| | | |
1575+-+-+-+-+-+-+-+-+
1576
Guido van Rossum7131f842007-02-09 20:13:25 +00001577>>> print(count, "solutions in all.")
Tim Petersbe4f0a72001-06-29 02:41:16 +0000157892 solutions in all.
unknown31569562001-07-04 22:11:22 +00001579
1580And run a Knight's Tour on a 10x10 board. Note that there are about
158120,000 solutions even on a 6x6 board, so don't dare run this to exhaustion.
1582
Tim Peters9a8c8e22001-07-13 09:12:12 +00001583>>> k = Knights(10, 10)
unknown31569562001-07-04 22:11:22 +00001584>>> LIMIT = 2
1585>>> count = 0
1586>>> for x in k.solve():
1587... count += 1
1588... if count <= LIMIT:
Guido van Rossum7131f842007-02-09 20:13:25 +00001589... print("Solution", count)
unknown31569562001-07-04 22:11:22 +00001590... k.printsolution(x)
1591... else:
1592... break
1593Solution 1
1594+---+---+---+---+---+---+---+---+---+---+
1595| 1| 58| 27| 34| 3| 40| 29| 10| 5| 8|
1596+---+---+---+---+---+---+---+---+---+---+
1597| 26| 35| 2| 57| 28| 33| 4| 7| 30| 11|
1598+---+---+---+---+---+---+---+---+---+---+
1599| 59|100| 73| 36| 41| 56| 39| 32| 9| 6|
1600+---+---+---+---+---+---+---+---+---+---+
1601| 74| 25| 60| 55| 72| 37| 42| 49| 12| 31|
1602+---+---+---+---+---+---+---+---+---+---+
1603| 61| 86| 99| 76| 63| 52| 47| 38| 43| 50|
1604+---+---+---+---+---+---+---+---+---+---+
1605| 24| 75| 62| 85| 54| 71| 64| 51| 48| 13|
1606+---+---+---+---+---+---+---+---+---+---+
1607| 87| 98| 91| 80| 77| 84| 53| 46| 65| 44|
1608+---+---+---+---+---+---+---+---+---+---+
1609| 90| 23| 88| 95| 70| 79| 68| 83| 14| 17|
1610+---+---+---+---+---+---+---+---+---+---+
1611| 97| 92| 21| 78| 81| 94| 19| 16| 45| 66|
1612+---+---+---+---+---+---+---+---+---+---+
1613| 22| 89| 96| 93| 20| 69| 82| 67| 18| 15|
1614+---+---+---+---+---+---+---+---+---+---+
1615Solution 2
1616+---+---+---+---+---+---+---+---+---+---+
1617| 1| 58| 27| 34| 3| 40| 29| 10| 5| 8|
1618+---+---+---+---+---+---+---+---+---+---+
1619| 26| 35| 2| 57| 28| 33| 4| 7| 30| 11|
1620+---+---+---+---+---+---+---+---+---+---+
1621| 59|100| 73| 36| 41| 56| 39| 32| 9| 6|
1622+---+---+---+---+---+---+---+---+---+---+
1623| 74| 25| 60| 55| 72| 37| 42| 49| 12| 31|
1624+---+---+---+---+---+---+---+---+---+---+
1625| 61| 86| 99| 76| 63| 52| 47| 38| 43| 50|
1626+---+---+---+---+---+---+---+---+---+---+
1627| 24| 75| 62| 85| 54| 71| 64| 51| 48| 13|
1628+---+---+---+---+---+---+---+---+---+---+
1629| 87| 98| 89| 80| 77| 84| 53| 46| 65| 44|
1630+---+---+---+---+---+---+---+---+---+---+
1631| 90| 23| 92| 95| 70| 79| 68| 83| 14| 17|
1632+---+---+---+---+---+---+---+---+---+---+
1633| 97| 88| 21| 78| 81| 94| 19| 16| 45| 66|
1634+---+---+---+---+---+---+---+---+---+---+
1635| 22| 91| 96| 93| 20| 69| 82| 67| 18| 15|
1636+---+---+---+---+---+---+---+---+---+---+
Tim Petersbe4f0a72001-06-29 02:41:16 +00001637"""
1638
Fred Drake56d12662002-08-09 18:37:10 +00001639weakref_tests = """\
1640Generators are weakly referencable:
1641
1642>>> import weakref
1643>>> def gen():
1644... yield 'foo!'
1645...
1646>>> wr = weakref.ref(gen)
1647>>> wr() is gen
1648True
1649>>> p = weakref.proxy(gen)
1650
1651Generator-iterators are weakly referencable as well:
1652
1653>>> gi = gen()
1654>>> wr = weakref.ref(gi)
1655>>> wr() is gi
1656True
1657>>> p = weakref.proxy(gi)
1658>>> list(p)
1659['foo!']
1660
1661"""
1662
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001663coroutine_tests = """\
1664Sending a value into a started generator:
1665
1666>>> def f():
Guido van Rossum7131f842007-02-09 20:13:25 +00001667... print((yield 1))
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001668... yield 2
1669>>> g = f()
Georg Brandla18af4e2007-04-21 15:47:16 +00001670>>> next(g)
Phillip J. Eby0d6615f2005-08-02 00:46:46 +000016711
1672>>> g.send(42)
167342
16742
1675
1676Sending a value into a new generator produces a TypeError:
1677
1678>>> f().send("foo")
1679Traceback (most recent call last):
1680...
1681TypeError: can't send non-None value to a just-started generator
1682
1683
1684Yield by itself yields None:
1685
1686>>> def f(): yield
1687>>> list(f())
1688[None]
1689
1690
1691
1692An obscene abuse of a yield expression within a generator expression:
1693
1694>>> list((yield 21) for i in range(4))
1695[21, None, 21, None, 21, None, 21, None]
1696
1697And a more sane, but still weird usage:
1698
1699>>> def f(): list(i for i in [(yield 26)])
1700>>> type(f())
Martin v. Löwis250ad612008-04-07 05:43:42 +00001701<class 'generator'>
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001702
1703
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001704A yield expression with augmented assignment.
1705
1706>>> def coroutine(seq):
1707... count = 0
1708... while count < 200:
1709... count += yield
1710... seq.append(count)
1711>>> seq = []
1712>>> c = coroutine(seq)
Georg Brandla18af4e2007-04-21 15:47:16 +00001713>>> next(c)
Guido van Rossum7131f842007-02-09 20:13:25 +00001714>>> print(seq)
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001715[]
1716>>> c.send(10)
Guido van Rossum7131f842007-02-09 20:13:25 +00001717>>> print(seq)
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001718[10]
1719>>> c.send(10)
Guido van Rossum7131f842007-02-09 20:13:25 +00001720>>> print(seq)
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001721[10, 20]
1722>>> c.send(10)
Guido van Rossum7131f842007-02-09 20:13:25 +00001723>>> print(seq)
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001724[10, 20, 30]
1725
1726
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001727Check some syntax errors for yield expressions:
1728
1729>>> f=lambda: (yield 1),(yield 2)
1730Traceback (most recent call last):
1731 ...
Guido van Rossum33d26892007-08-05 15:29:28 +00001732SyntaxError: 'yield' outside function
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001733
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001734>>> def f(): x = yield = y
1735Traceback (most recent call last):
1736 ...
Guido van Rossum33d26892007-08-05 15:29:28 +00001737SyntaxError: assignment to yield expression not possible
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001738
1739>>> def f(): (yield bar) = y
1740Traceback (most recent call last):
1741 ...
Guido van Rossum33d26892007-08-05 15:29:28 +00001742SyntaxError: can't assign to yield expression
Thomas Wouters0e3f5912006-08-11 14:57:12 +00001743
1744>>> def f(): (yield bar) += y
1745Traceback (most recent call last):
1746 ...
Benjamin Peterson87c8d872009-06-11 22:54:11 +00001747SyntaxError: can't assign to yield expression
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001748
1749
1750Now check some throw() conditions:
1751
1752>>> def f():
1753... while True:
1754... try:
Guido van Rossum7131f842007-02-09 20:13:25 +00001755... print((yield))
Guido van Rossumb940e112007-01-10 16:19:56 +00001756... except ValueError as v:
Guido van Rossumc420b2f2007-02-09 22:09:01 +00001757... print("caught ValueError (%s)" % (v))
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001758>>> import sys
1759>>> g = f()
Georg Brandla18af4e2007-04-21 15:47:16 +00001760>>> next(g)
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001761
1762>>> g.throw(ValueError) # type only
1763caught ValueError ()
1764
1765>>> g.throw(ValueError("xyz")) # value only
1766caught ValueError (xyz)
1767
1768>>> g.throw(ValueError, ValueError(1)) # value+matching type
1769caught ValueError (1)
1770
1771>>> g.throw(ValueError, TypeError(1)) # mismatched type, rewrapped
1772caught ValueError (1)
1773
Thomas Wouters49fd7fa2006-04-21 10:40:58 +00001774>>> g.throw(ValueError, ValueError(1), None) # explicit None traceback
1775caught ValueError (1)
1776
Tim Peterse9fe7e02005-08-07 03:04:58 +00001777>>> g.throw(ValueError(1), "foo") # bad args
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001778Traceback (most recent call last):
1779 ...
1780TypeError: instance exception may not have a separate value
1781
Tim Peterse9fe7e02005-08-07 03:04:58 +00001782>>> g.throw(ValueError, "foo", 23) # bad args
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001783Traceback (most recent call last):
1784 ...
1785TypeError: throw() third argument must be a traceback object
1786
Guido van Rossumbf12cdb2006-08-17 20:24:18 +00001787>>> g.throw("abc")
1788Traceback (most recent call last):
1789 ...
1790TypeError: exceptions must be classes or instances deriving from BaseException, not str
1791
1792>>> g.throw(0)
1793Traceback (most recent call last):
1794 ...
1795TypeError: exceptions must be classes or instances deriving from BaseException, not int
1796
1797>>> g.throw(list)
1798Traceback (most recent call last):
1799 ...
1800TypeError: exceptions must be classes or instances deriving from BaseException, not type
1801
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001802>>> def throw(g,exc):
1803... try:
1804... raise exc
1805... except:
1806... g.throw(*sys.exc_info())
Tim Peterse9fe7e02005-08-07 03:04:58 +00001807>>> throw(g,ValueError) # do it with traceback included
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001808caught ValueError ()
1809
1810>>> g.send(1)
18111
1812
Tim Peterse9fe7e02005-08-07 03:04:58 +00001813>>> throw(g,TypeError) # terminate the generator
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001814Traceback (most recent call last):
1815 ...
1816TypeError
1817
Guido van Rossum7131f842007-02-09 20:13:25 +00001818>>> print(g.gi_frame)
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001819None
1820
1821>>> g.send(2)
1822Traceback (most recent call last):
1823 ...
1824StopIteration
1825
Tim Peterse9fe7e02005-08-07 03:04:58 +00001826>>> g.throw(ValueError,6) # throw on closed generator
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001827Traceback (most recent call last):
1828 ...
1829ValueError: 6
1830
Tim Peterse9fe7e02005-08-07 03:04:58 +00001831>>> f().throw(ValueError,7) # throw on just-opened generator
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001832Traceback (most recent call last):
1833 ...
1834ValueError: 7
1835
Antoine Pitrou551ba202011-10-18 16:40:50 +02001836Plain "raise" inside a generator should preserve the traceback (#13188).
1837The traceback should have 3 levels:
1838- g.throw()
1839- f()
1840- 1/0
1841
1842>>> def f():
1843... try:
1844... yield
1845... except:
1846... raise
1847>>> g = f()
1848>>> try:
1849... 1/0
1850... except ZeroDivisionError as v:
1851... try:
1852... g.throw(v)
1853... except Exception as w:
1854... tb = w.__traceback__
1855>>> levels = 0
1856>>> while tb:
1857... levels += 1
1858... tb = tb.tb_next
1859>>> levels
18603
1861
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001862Now let's try closing a generator:
1863
1864>>> def f():
1865... try: yield
1866... except GeneratorExit:
Guido van Rossum7131f842007-02-09 20:13:25 +00001867... print("exiting")
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001868
1869>>> g = f()
Georg Brandla18af4e2007-04-21 15:47:16 +00001870>>> next(g)
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001871>>> g.close()
1872exiting
1873>>> g.close() # should be no-op now
1874
1875>>> f().close() # close on just-opened generator should be fine
1876
Tim Peterse9fe7e02005-08-07 03:04:58 +00001877>>> def f(): yield # an even simpler generator
1878>>> f().close() # close before opening
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001879>>> g = f()
Georg Brandla18af4e2007-04-21 15:47:16 +00001880>>> next(g)
Tim Peterse9fe7e02005-08-07 03:04:58 +00001881>>> g.close() # close normally
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001882
1883And finalization:
1884
1885>>> def f():
1886... try: yield
1887... finally:
Guido van Rossum7131f842007-02-09 20:13:25 +00001888... print("exiting")
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001889
1890>>> g = f()
Georg Brandla18af4e2007-04-21 15:47:16 +00001891>>> next(g)
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001892>>> del g
1893exiting
1894
1895
Christian Heimescbf3b5c2007-12-03 21:02:03 +00001896GeneratorExit is not caught by except Exception:
1897
1898>>> def f():
1899... try: yield
1900... except Exception:
1901... print('except')
1902... finally:
1903... print('finally')
1904
1905>>> g = f()
1906>>> next(g)
1907>>> del g
1908finally
1909
1910
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001911Now let's try some ill-behaved generators:
1912
1913>>> def f():
1914... try: yield
1915... except GeneratorExit:
1916... yield "foo!"
1917>>> g = f()
Georg Brandla18af4e2007-04-21 15:47:16 +00001918>>> next(g)
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001919>>> g.close()
1920Traceback (most recent call last):
1921 ...
1922RuntimeError: generator ignored GeneratorExit
1923>>> g.close()
1924
1925
1926Our ill-behaved code should be invoked during GC:
1927
Guido van Rossum34d19282007-08-09 01:03:29 +00001928>>> import sys, io
1929>>> old, sys.stderr = sys.stderr, io.StringIO()
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001930>>> g = f()
Georg Brandla18af4e2007-04-21 15:47:16 +00001931>>> next(g)
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001932>>> del g
Andrew Svetlov76bcff22012-11-03 15:56:05 +02001933>>> "RuntimeError: generator ignored GeneratorExit" in sys.stderr.getvalue()
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001934True
1935>>> sys.stderr = old
1936
1937
1938And errors thrown during closing should propagate:
1939
1940>>> def f():
1941... try: yield
1942... except GeneratorExit:
1943... raise TypeError("fie!")
1944>>> g = f()
Georg Brandla18af4e2007-04-21 15:47:16 +00001945>>> next(g)
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001946>>> g.close()
1947Traceback (most recent call last):
1948 ...
1949TypeError: fie!
1950
1951
1952Ensure that various yield expression constructs make their
1953enclosing function a generator:
1954
1955>>> def f(): x += yield
1956>>> type(f())
Martin v. Löwis250ad612008-04-07 05:43:42 +00001957<class 'generator'>
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001958
1959>>> def f(): x = yield
1960>>> type(f())
Martin v. Löwis250ad612008-04-07 05:43:42 +00001961<class 'generator'>
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001962
1963>>> def f(): lambda x=(yield): 1
1964>>> type(f())
Martin v. Löwis250ad612008-04-07 05:43:42 +00001965<class 'generator'>
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001966
1967>>> def f(): x=(i for i in (yield) if (yield))
1968>>> type(f())
Martin v. Löwis250ad612008-04-07 05:43:42 +00001969<class 'generator'>
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001970
1971>>> def f(d): d[(yield "a")] = d[(yield "b")] = 27
1972>>> data = [1,2]
1973>>> g = f(data)
1974>>> type(g)
Martin v. Löwis250ad612008-04-07 05:43:42 +00001975<class 'generator'>
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00001976>>> g.send(None)
1977'a'
1978>>> data
1979[1, 2]
1980>>> g.send(0)
1981'b'
1982>>> data
1983[27, 2]
1984>>> try: g.send(1)
1985... except StopIteration: pass
1986>>> data
1987[27, 27]
1988
1989"""
1990
Thomas Wouters49fd7fa2006-04-21 10:40:58 +00001991refleaks_tests = """
1992Prior to adding cycle-GC support to itertools.tee, this code would leak
1993references. We add it to the standard suite so the routine refleak-tests
1994would trigger if it starts being uncleanable again.
1995
1996>>> import itertools
1997>>> def leak():
1998... class gen:
1999... def __iter__(self):
2000... return self
Georg Brandla18af4e2007-04-21 15:47:16 +00002001... def __next__(self):
Thomas Wouters49fd7fa2006-04-21 10:40:58 +00002002... return self.item
2003... g = gen()
2004... head, tail = itertools.tee(g)
2005... g.item = head
2006... return head
2007>>> it = leak()
2008
2009Make sure to also test the involvement of the tee-internal teedataobject,
2010which stores returned items.
2011
Georg Brandla18af4e2007-04-21 15:47:16 +00002012>>> item = next(it)
Thomas Wouters49fd7fa2006-04-21 10:40:58 +00002013
2014
2015
2016This test leaked at one point due to generator finalization/destruction.
2017It was copied from Lib/test/leakers/test_generator_cycle.py before the file
2018was removed.
2019
2020>>> def leak():
2021... def gen():
2022... while True:
2023... yield g
2024... g = gen()
2025
2026>>> leak()
2027
2028
2029
2030This test isn't really generator related, but rather exception-in-cleanup
2031related. The coroutine tests (above) just happen to cause an exception in
2032the generator's __del__ (tp_del) method. We can also test for this
2033explicitly, without generators. We do have to redirect stderr to avoid
2034printing warnings and to doublecheck that we actually tested what we wanted
2035to test.
2036
Guido van Rossum34d19282007-08-09 01:03:29 +00002037>>> import sys, io
Thomas Wouters49fd7fa2006-04-21 10:40:58 +00002038>>> old = sys.stderr
2039>>> try:
Guido van Rossum34d19282007-08-09 01:03:29 +00002040... sys.stderr = io.StringIO()
Thomas Wouters49fd7fa2006-04-21 10:40:58 +00002041... class Leaker:
2042... def __del__(self):
Andrew Svetlov76bcff22012-11-03 15:56:05 +02002043... def invoke(message):
2044... raise RuntimeError(message)
2045... invoke("test")
Thomas Wouters49fd7fa2006-04-21 10:40:58 +00002046...
2047... l = Leaker()
2048... del l
2049... err = sys.stderr.getvalue().strip()
Andrew Svetlov76bcff22012-11-03 15:56:05 +02002050... "Exception ignored in" in err
2051... "RuntimeError: test" in err
2052... "Traceback" in err
2053... "in invoke" in err
Thomas Wouters49fd7fa2006-04-21 10:40:58 +00002054... finally:
2055... sys.stderr = old
2056True
2057True
Andrew Svetlov76bcff22012-11-03 15:56:05 +02002058True
2059True
Thomas Wouters49fd7fa2006-04-21 10:40:58 +00002060
2061
2062These refleak tests should perhaps be in a testfile of their own,
2063test_generators just happened to be the test that drew these out.
2064
2065"""
2066
Tim Petersf6ed0742001-06-27 07:17:57 +00002067__test__ = {"tut": tutorial_tests,
2068 "pep": pep_tests,
2069 "email": email_tests,
2070 "fun": fun_tests,
Tim Petersbe4f0a72001-06-29 02:41:16 +00002071 "syntax": syntax_tests,
Fred Drake56d12662002-08-09 18:37:10 +00002072 "conjoin": conjoin_tests,
2073 "weakref": weakref_tests,
Phillip J. Eby0d6615f2005-08-02 00:46:46 +00002074 "coroutine": coroutine_tests,
Thomas Wouters49fd7fa2006-04-21 10:40:58 +00002075 "refleaks": refleaks_tests,
Fred Drake56d12662002-08-09 18:37:10 +00002076 }
Tim Peters1def3512001-06-23 20:27:04 +00002077
2078# Magic test name that regrtest.py invokes *after* importing this module.
2079# This worms around a bootstrap problem.
2080# Note that doctest and regrtest both look in sys.argv for a "-v" argument,
2081# so this works as expected in both ways of running regrtest.
Tim Petersa0a62222001-09-09 06:12:01 +00002082def test_main(verbose=None):
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002083 from test import support, test_generators
Antoine Pitrou796564c2013-07-30 19:59:21 +02002084 support.run_unittest(__name__)
Benjamin Petersonee8712c2008-05-20 21:35:26 +00002085 support.run_doctest(test_generators, verbose)
Tim Peters1def3512001-06-23 20:27:04 +00002086
2087# This part isn't needed for regrtest, but for running the test directly.
2088if __name__ == "__main__":
Tim Petersa0a62222001-09-09 06:12:01 +00002089 test_main(1)