blob: ee7d34878c66c3c66b62d4f4b18afac4dfd62eee [file] [log] [blame]
Tim Petersbe4f0a72001-06-29 02:41:16 +00001from __future__ import nested_scopes
2
Tim Peters6ba5f792001-06-23 20:45:43 +00003tutorial_tests = """
Tim Peters1def3512001-06-23 20:27:04 +00004Let's try a simple generator:
5
6 >>> def f():
7 ... yield 1
8 ... yield 2
9
Tim Petersb9e9ff12001-06-24 03:44:52 +000010 >>> for i in f():
11 ... print i
12 1
13 2
Tim Peters1def3512001-06-23 20:27:04 +000014 >>> g = f()
15 >>> g.next()
16 1
17 >>> g.next()
18 2
Tim Petersea2e97a2001-06-24 07:10:02 +000019
Tim Peters2106ef02001-06-25 01:30:12 +000020"Falling off the end" stops the generator:
Tim Petersea2e97a2001-06-24 07:10:02 +000021
Tim Peters1def3512001-06-23 20:27:04 +000022 >>> g.next()
23 Traceback (most recent call last):
24 File "<stdin>", line 1, in ?
25 File "<stdin>", line 2, in g
26 StopIteration
27
Tim Petersea2e97a2001-06-24 07:10:02 +000028"return" also stops the generator:
Tim Peters1def3512001-06-23 20:27:04 +000029
30 >>> def f():
31 ... yield 1
32 ... return
33 ... yield 2 # never reached
34 ...
35 >>> g = f()
36 >>> g.next()
37 1
38 >>> g.next()
39 Traceback (most recent call last):
40 File "<stdin>", line 1, in ?
41 File "<stdin>", line 3, in f
42 StopIteration
43 >>> g.next() # once stopped, can't be resumed
44 Traceback (most recent call last):
45 File "<stdin>", line 1, in ?
46 StopIteration
47
48"raise StopIteration" stops the generator too:
49
50 >>> def f():
51 ... yield 1
52 ... return
53 ... yield 2 # never reached
54 ...
55 >>> g = f()
56 >>> g.next()
57 1
58 >>> g.next()
59 Traceback (most recent call last):
60 File "<stdin>", line 1, in ?
61 StopIteration
62 >>> g.next()
63 Traceback (most recent call last):
64 File "<stdin>", line 1, in ?
65 StopIteration
66
67However, they are not exactly equivalent:
68
69 >>> def g1():
70 ... try:
71 ... return
72 ... except:
73 ... yield 1
74 ...
75 >>> list(g1())
76 []
77
78 >>> def g2():
79 ... try:
80 ... raise StopIteration
81 ... except:
82 ... yield 42
83 >>> print list(g2())
84 [42]
85
86This may be surprising at first:
87
88 >>> def g3():
89 ... try:
90 ... return
91 ... finally:
92 ... yield 1
93 ...
94 >>> list(g3())
95 [1]
96
97Let's create an alternate range() function implemented as a generator:
98
99 >>> def yrange(n):
100 ... for i in range(n):
101 ... yield i
102 ...
103 >>> list(yrange(5))
104 [0, 1, 2, 3, 4]
105
106Generators always return to the most recent caller:
107
108 >>> def creator():
109 ... r = yrange(5)
110 ... print "creator", r.next()
111 ... return r
112 ...
113 >>> def caller():
114 ... r = creator()
115 ... for i in r:
116 ... print "caller", i
117 ...
118 >>> caller()
119 creator 0
120 caller 1
121 caller 2
122 caller 3
123 caller 4
124
125Generators can call other generators:
126
127 >>> def zrange(n):
128 ... for i in yrange(n):
129 ... yield i
130 ...
131 >>> list(zrange(5))
132 [0, 1, 2, 3, 4]
133
134"""
135
Tim Peters6ba5f792001-06-23 20:45:43 +0000136# The examples from PEP 255.
137
138pep_tests = """
139
140Specification: Return
141
142 Note that return isn't always equivalent to raising StopIteration: the
143 difference lies in how enclosing try/except constructs are treated.
144 For example,
145
146 >>> def f1():
147 ... try:
148 ... return
149 ... except:
150 ... yield 1
151 >>> print list(f1())
152 []
153
154 because, as in any function, return simply exits, but
155
156 >>> def f2():
157 ... try:
158 ... raise StopIteration
159 ... except:
160 ... yield 42
161 >>> print list(f2())
162 [42]
163
164 because StopIteration is captured by a bare "except", as is any
165 exception.
166
167Specification: Generators and Exception Propagation
168
169 >>> def f():
170 ... return 1/0
171 >>> def g():
172 ... yield f() # the zero division exception propagates
173 ... yield 42 # and we'll never get here
174 >>> k = g()
175 >>> k.next()
176 Traceback (most recent call last):
177 File "<stdin>", line 1, in ?
178 File "<stdin>", line 2, in g
179 File "<stdin>", line 2, in f
180 ZeroDivisionError: integer division or modulo by zero
181 >>> k.next() # and the generator cannot be resumed
182 Traceback (most recent call last):
183 File "<stdin>", line 1, in ?
184 StopIteration
185 >>>
186
187Specification: Try/Except/Finally
188
189 >>> def f():
190 ... try:
191 ... yield 1
192 ... try:
193 ... yield 2
194 ... 1/0
195 ... yield 3 # never get here
196 ... except ZeroDivisionError:
197 ... yield 4
198 ... yield 5
199 ... raise
200 ... except:
201 ... yield 6
202 ... yield 7 # the "raise" above stops this
203 ... except:
204 ... yield 8
205 ... yield 9
206 ... try:
207 ... x = 12
208 ... finally:
209 ... yield 10
210 ... yield 11
211 >>> print list(f())
212 [1, 2, 4, 5, 8, 9, 10, 11]
213 >>>
214
Tim Peters6ba5f792001-06-23 20:45:43 +0000215Guido's binary tree example.
216
217 >>> # A binary tree class.
218 >>> class Tree:
219 ...
220 ... def __init__(self, label, left=None, right=None):
221 ... self.label = label
222 ... self.left = left
223 ... self.right = right
224 ...
225 ... def __repr__(self, level=0, indent=" "):
226 ... s = level*indent + `self.label`
227 ... if self.left:
228 ... s = s + "\\n" + self.left.__repr__(level+1, indent)
229 ... if self.right:
230 ... s = s + "\\n" + self.right.__repr__(level+1, indent)
231 ... return s
232 ...
233 ... def __iter__(self):
234 ... return inorder(self)
235
236 >>> # Create a Tree from a list.
237 >>> def tree(list):
238 ... n = len(list)
239 ... if n == 0:
240 ... return []
241 ... i = n / 2
242 ... return Tree(list[i], tree(list[:i]), tree(list[i+1:]))
243
244 >>> # Show it off: create a tree.
245 >>> t = tree("ABCDEFGHIJKLMNOPQRSTUVWXYZ")
246
247 >>> # A recursive generator that generates Tree leaves in in-order.
248 >>> def inorder(t):
249 ... if t:
250 ... for x in inorder(t.left):
251 ... yield x
252 ... yield t.label
253 ... for x in inorder(t.right):
254 ... yield x
255
256 >>> # Show it off: create a tree.
257 ... t = tree("ABCDEFGHIJKLMNOPQRSTUVWXYZ")
258 ... # Print the nodes of the tree in in-order.
259 ... for x in t:
260 ... print x,
261 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
262
263 >>> # A non-recursive generator.
264 >>> def inorder(node):
265 ... stack = []
266 ... while node:
267 ... while node.left:
268 ... stack.append(node)
269 ... node = node.left
270 ... yield node.label
271 ... while not node.right:
272 ... try:
273 ... node = stack.pop()
274 ... except IndexError:
275 ... return
276 ... yield node.label
277 ... node = node.right
278
279 >>> # Exercise the non-recursive generator.
280 >>> for x in t:
281 ... print x,
282 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
283
284"""
285
Tim Petersb2bc6a92001-06-24 10:14:27 +0000286# Examples from Iterator-List and Python-Dev and c.l.py.
Tim Peters6ba5f792001-06-23 20:45:43 +0000287
288email_tests = """
289
290The difference between yielding None and returning it.
291
292>>> def g():
293... for i in range(3):
294... yield None
295... yield None
296... return
297>>> list(g())
298[None, None, None, None]
299
300Ensure that explicitly raising StopIteration acts like any other exception
301in try/except, not like a return.
302
303>>> def g():
304... yield 1
305... try:
306... raise StopIteration
307... except:
308... yield 2
309... yield 3
310>>> list(g())
311[1, 2, 3]
Tim Petersb9e9ff12001-06-24 03:44:52 +0000312
313A generator can't be resumed while it's already running.
314
315>>> def g():
316... i = me.next()
317... yield i
318>>> me = g()
319>>> me.next()
320Traceback (most recent call last):
321 ...
322 File "<string>", line 2, in g
323ValueError: generator already executing
Tim Petersb2bc6a92001-06-24 10:14:27 +0000324
325Next one was posted to c.l.py.
326
327>>> def gcomb(x, k):
328... "Generate all combinations of k elements from list x."
329...
330... if k > len(x):
331... return
332... if k == 0:
333... yield []
334... else:
335... first, rest = x[0], x[1:]
336... # A combination does or doesn't contain first.
337... # If it does, the remainder is a k-1 comb of rest.
338... for c in gcomb(rest, k-1):
339... c.insert(0, first)
340... yield c
341... # If it doesn't contain first, it's a k comb of rest.
342... for c in gcomb(rest, k):
343... yield c
344
345>>> seq = range(1, 5)
346>>> for k in range(len(seq) + 2):
347... print "%d-combs of %s:" % (k, seq)
348... for c in gcomb(seq, k):
349... print " ", c
3500-combs of [1, 2, 3, 4]:
351 []
3521-combs of [1, 2, 3, 4]:
353 [1]
354 [2]
355 [3]
356 [4]
3572-combs of [1, 2, 3, 4]:
358 [1, 2]
359 [1, 3]
360 [1, 4]
361 [2, 3]
362 [2, 4]
363 [3, 4]
3643-combs of [1, 2, 3, 4]:
365 [1, 2, 3]
366 [1, 2, 4]
367 [1, 3, 4]
368 [2, 3, 4]
3694-combs of [1, 2, 3, 4]:
370 [1, 2, 3, 4]
3715-combs of [1, 2, 3, 4]:
Tim Peters3e7b1a02001-06-25 19:46:25 +0000372
Tim Peterse77f2e22001-06-26 22:24:51 +0000373From the Iterators list, about the types of these things.
Tim Peters3e7b1a02001-06-25 19:46:25 +0000374
375>>> def g():
376... yield 1
377...
378>>> type(g)
379<type 'function'>
380>>> i = g()
381>>> type(i)
382<type 'generator'>
383>>> dir(i)
Tim Peterse77f2e22001-06-26 22:24:51 +0000384['gi_frame', 'gi_running', 'next']
Tim Peters3e7b1a02001-06-25 19:46:25 +0000385>>> print i.next.__doc__
386next() -- get the next value, or raise StopIteration
387>>> iter(i) is i
3881
389>>> import types
390>>> isinstance(i, types.GeneratorType)
3911
Tim Peterse77f2e22001-06-26 22:24:51 +0000392
393And more, added later.
394
395>>> i.gi_running
3960
397>>> type(i.gi_frame)
398<type 'frame'>
399>>> i.gi_running = 42
400Traceback (most recent call last):
401 ...
402TypeError: object has read-only attributes
403>>> def g():
404... yield me.gi_running
405>>> me = g()
406>>> me.gi_running
4070
408>>> me.next()
4091
410>>> me.gi_running
4110
Tim Peters35302662001-07-02 01:38:33 +0000412
413A clever union-find implementation from c.l.py, due to David Eppstein.
414Sent: Friday, June 29, 2001 12:16 PM
415To: python-list@python.org
416Subject: Re: PEP 255: Simple Generators
417
418>>> class disjointSet:
419... def __init__(self, name):
420... self.name = name
421... self.parent = None
422... self.generator = self.generate()
423...
424... def generate(self):
425... while not self.parent:
426... yield self
427... for x in self.parent.generator:
428... yield x
429...
430... def find(self):
431... return self.generator.next()
432...
433... def union(self, parent):
434... if self.parent:
435... raise ValueError("Sorry, I'm not a root!")
436... self.parent = parent
437...
438... def __str__(self):
439... return self.name
440...
441... def clear(self):
442... self.__dict__.clear()
443
444>>> names = "ABCDEFGHIJKLM"
445>>> sets = [disjointSet(name) for name in names]
446>>> roots = sets[:]
447
448>>> import random
449>>> random.seed(42)
450>>> while 1:
451... for s in sets:
452... print "%s->%s" % (s, s.find()),
453... print
454... if len(roots) > 1:
455... s1 = random.choice(roots)
456... roots.remove(s1)
457... s2 = random.choice(roots)
458... s1.union(s2)
459... print "merged", s1, "into", s2
460... else:
461... break
462A->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
463merged D into G
464A->A B->B C->C D->G E->E F->F G->G H->H I->I J->J K->K L->L M->M
465merged C into F
466A->A B->B C->F D->G E->E F->F G->G H->H I->I J->J K->K L->L M->M
467merged L into A
468A->A B->B C->F D->G E->E F->F G->G H->H I->I J->J K->K L->A M->M
469merged H into E
470A->A B->B C->F D->G E->E F->F G->G H->E I->I J->J K->K L->A M->M
471merged B into E
472A->A B->E C->F D->G E->E F->F G->G H->E I->I J->J K->K L->A M->M
473merged J into G
474A->A B->E C->F D->G E->E F->F G->G H->E I->I J->G K->K L->A M->M
475merged E into G
476A->A B->G C->F D->G E->G F->F G->G H->G I->I J->G K->K L->A M->M
477merged M into G
478A->A B->G C->F D->G E->G F->F G->G H->G I->I J->G K->K L->A M->G
479merged I into K
480A->A B->G C->F D->G E->G F->F G->G H->G I->K J->G K->K L->A M->G
481merged K into A
482A->A B->G C->F D->G E->G F->F G->G H->G I->A J->G K->A L->A M->G
483merged F into A
484A->A B->G C->A D->G E->G F->A G->G H->G I->A J->G K->A L->A M->G
485merged A into G
486A->G B->G C->G D->G E->G F->G G->G H->G I->G J->G K->G L->G M->G
487>>> for s in sets: # XXX memory leak without this
488... s.clear()
Tim Peters6ba5f792001-06-23 20:45:43 +0000489"""
490
Tim Peters0f9da0a2001-06-23 21:01:47 +0000491# Fun tests (for sufficiently warped notions of "fun").
492
493fun_tests = """
494
495Build up to a recursive Sieve of Eratosthenes generator.
496
497>>> def firstn(g, n):
498... return [g.next() for i in range(n)]
499
500>>> def intsfrom(i):
501... while 1:
502... yield i
503... i += 1
504
505>>> firstn(intsfrom(5), 7)
506[5, 6, 7, 8, 9, 10, 11]
507
508>>> def exclude_multiples(n, ints):
509... for i in ints:
510... if i % n:
511... yield i
512
513>>> firstn(exclude_multiples(3, intsfrom(1)), 6)
514[1, 2, 4, 5, 7, 8]
515
516>>> def sieve(ints):
517... prime = ints.next()
518... yield prime
519... not_divisible_by_prime = exclude_multiples(prime, ints)
520... for p in sieve(not_divisible_by_prime):
521... yield p
522
523>>> primes = sieve(intsfrom(2))
524>>> firstn(primes, 20)
525[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 +0000526
Tim Petersf6ed0742001-06-27 07:17:57 +0000527
Tim Petersb9e9ff12001-06-24 03:44:52 +0000528Another famous problem: generate all integers of the form
529 2**i * 3**j * 5**k
530in increasing order, where i,j,k >= 0. Trickier than it may look at first!
531Try writing it without generators, and correctly, and without generating
5323 internal results for each result output.
533
534>>> def times(n, g):
535... for i in g:
536... yield n * i
537>>> firstn(times(10, intsfrom(1)), 10)
538[10, 20, 30, 40, 50, 60, 70, 80, 90, 100]
539
540>>> def merge(g, h):
541... ng = g.next()
542... nh = h.next()
543... while 1:
544... if ng < nh:
545... yield ng
546... ng = g.next()
547... elif ng > nh:
548... yield nh
549... nh = h.next()
550... else:
551... yield ng
552... ng = g.next()
553... nh = h.next()
554
Tim Petersf6ed0742001-06-27 07:17:57 +0000555The following works, but is doing a whale of a lot of redundant work --
556it's not clear how to get the internal uses of m235 to share a single
557generator. Note that me_times2 (etc) each need to see every element in the
558result sequence. So this is an example where lazy lists are more natural
559(you can look at the head of a lazy list any number of times).
Tim Petersb9e9ff12001-06-24 03:44:52 +0000560
561>>> def m235():
562... yield 1
563... me_times2 = times(2, m235())
564... me_times3 = times(3, m235())
565... me_times5 = times(5, m235())
566... for i in merge(merge(me_times2,
567... me_times3),
568... me_times5):
569... yield i
570
Tim Petersf6ed0742001-06-27 07:17:57 +0000571Don't print "too many" of these -- the implementation above is extremely
572inefficient: each call of m235() leads to 3 recursive calls, and in
573turn each of those 3 more, and so on, and so on, until we've descended
574enough levels to satisfy the print stmts. Very odd: when I printed 5
575lines of results below, this managed to screw up Win98's malloc in "the
576usual" way, i.e. the heap grew over 4Mb so Win98 started fragmenting
577address space, and it *looked* like a very slow leak.
578
Tim Petersb9e9ff12001-06-24 03:44:52 +0000579>>> result = m235()
Tim Petersf6ed0742001-06-27 07:17:57 +0000580>>> for i in range(3):
Tim Petersb9e9ff12001-06-24 03:44:52 +0000581... print firstn(result, 15)
582[1, 2, 3, 4, 5, 6, 8, 9, 10, 12, 15, 16, 18, 20, 24]
583[25, 27, 30, 32, 36, 40, 45, 48, 50, 54, 60, 64, 72, 75, 80]
584[81, 90, 96, 100, 108, 120, 125, 128, 135, 144, 150, 160, 162, 180, 192]
Tim Petersee309272001-06-24 05:47:06 +0000585
586Heh. Here's one way to get a shared list, complete with an excruciating
587namespace renaming trick. The *pretty* part is that the times() and merge()
588functions can be reused as-is, because they only assume their stream
589arguments are iterable -- a LazyList is the same as a generator to times().
590
591>>> class LazyList:
592... def __init__(self, g):
593... self.sofar = []
594... self.fetch = g.next
595...
596... def __getitem__(self, i):
597... sofar, fetch = self.sofar, self.fetch
598... while i >= len(sofar):
599... sofar.append(fetch())
600... return sofar[i]
Tim Petersf6ed0742001-06-27 07:17:57 +0000601...
602... def clear(self):
603... self.__dict__.clear()
Tim Petersee309272001-06-24 05:47:06 +0000604
605>>> def m235():
606... yield 1
Tim Petersea2e97a2001-06-24 07:10:02 +0000607... # Gack: m235 below actually refers to a LazyList.
Tim Petersee309272001-06-24 05:47:06 +0000608... me_times2 = times(2, m235)
609... me_times3 = times(3, m235)
610... me_times5 = times(5, m235)
611... for i in merge(merge(me_times2,
612... me_times3),
613... me_times5):
614... yield i
615
Tim Petersf6ed0742001-06-27 07:17:57 +0000616Print as many of these as you like -- *this* implementation is memory-
617efficient. XXX Except that it leaks unless you clear the dict!
618
Tim Petersee309272001-06-24 05:47:06 +0000619>>> m235 = LazyList(m235())
620>>> for i in range(5):
621... print [m235[j] for j in range(15*i, 15*(i+1))]
622[1, 2, 3, 4, 5, 6, 8, 9, 10, 12, 15, 16, 18, 20, 24]
623[25, 27, 30, 32, 36, 40, 45, 48, 50, 54, 60, 64, 72, 75, 80]
624[81, 90, 96, 100, 108, 120, 125, 128, 135, 144, 150, 160, 162, 180, 192]
625[200, 216, 225, 240, 243, 250, 256, 270, 288, 300, 320, 324, 360, 375, 384]
626[400, 405, 432, 450, 480, 486, 500, 512, 540, 576, 600, 625, 640, 648, 675]
Tim Petersf6ed0742001-06-27 07:17:57 +0000627
628>>> m235.clear() # XXX memory leak without this
629
630
631Ye olde Fibonacci generator, LazyList style.
632
633>>> def fibgen(a, b):
634...
635... def sum(g, h):
636... while 1:
637... yield g.next() + h.next()
638...
639... def tail(g):
640... g.next() # throw first away
641... for x in g:
642... yield x
643...
644... yield a
645... yield b
646... for s in sum(iter(fib),
647... tail(iter(fib))):
648... yield s
649
650>>> fib = LazyList(fibgen(1, 2))
651>>> firstn(iter(fib), 17)
652[1, 2, 3, 5, 8, 13, 21, 34, 55, 89, 144, 233, 377, 610, 987, 1597, 2584]
653
654>>> fib.clear() # XXX memory leak without this
Tim Peters0f9da0a2001-06-23 21:01:47 +0000655"""
656
Tim Petersb6c3cea2001-06-26 03:36:28 +0000657# syntax_tests mostly provokes SyntaxErrors. Also fiddling with #if 0
658# hackery.
Tim Petersee309272001-06-24 05:47:06 +0000659
Tim Petersea2e97a2001-06-24 07:10:02 +0000660syntax_tests = """
661
662>>> def f():
663... return 22
664... yield 1
665Traceback (most recent call last):
666 ...
667SyntaxError: 'return' with argument inside generator (<string>, line 2)
668
669>>> def f():
670... yield 1
671... return 22
672Traceback (most recent call last):
673 ...
674SyntaxError: 'return' with argument inside generator (<string>, line 3)
675
676"return None" is not the same as "return" in a generator:
677
678>>> def f():
679... yield 1
680... return None
681Traceback (most recent call last):
682 ...
683SyntaxError: 'return' with argument inside generator (<string>, line 3)
684
685This one is fine:
686
687>>> def f():
688... yield 1
689... return
690
691>>> def f():
692... try:
693... yield 1
694... finally:
695... pass
696Traceback (most recent call last):
697 ...
698SyntaxError: 'yield' not allowed in a 'try' block with a 'finally' clause (<string>, line 3)
699
700>>> def f():
701... try:
702... try:
703... 1/0
704... except ZeroDivisionError:
705... yield 666 # bad because *outer* try has finally
706... except:
707... pass
708... finally:
709... pass
710Traceback (most recent call last):
711 ...
712SyntaxError: 'yield' not allowed in a 'try' block with a 'finally' clause (<string>, line 6)
713
714But this is fine:
715
716>>> def f():
717... try:
718... try:
719... yield 12
720... 1/0
721... except ZeroDivisionError:
722... yield 666
723... except:
724... try:
725... x = 12
726... finally:
727... yield 12
728... except:
729... return
730>>> list(f())
731[12, 666]
Tim Petersb6c3cea2001-06-26 03:36:28 +0000732
733>>> def f():
Tim Peters08a898f2001-06-28 01:52:22 +0000734... yield
735Traceback (most recent call last):
736SyntaxError: invalid syntax
737
738>>> def f():
739... if 0:
740... yield
741Traceback (most recent call last):
742SyntaxError: invalid syntax
743
744>>> def f():
Tim Petersb6c3cea2001-06-26 03:36:28 +0000745... if 0:
746... yield 1
747>>> type(f())
748<type 'generator'>
749
750>>> def f():
751... if "":
752... yield None
753>>> type(f())
754<type 'generator'>
755
756>>> def f():
757... return
758... try:
759... if x==4:
760... pass
761... elif 0:
762... try:
763... 1/0
764... except SyntaxError:
765... pass
766... else:
767... if 0:
768... while 12:
769... x += 1
770... yield 2 # don't blink
771... f(a, b, c, d, e)
772... else:
773... pass
774... except:
775... x = 1
776... return
777>>> type(f())
778<type 'generator'>
779
780>>> def f():
781... if 0:
782... def g():
783... yield 1
784...
785>>> type(f())
786<type 'None'>
787
788>>> def f():
789... if 0:
790... class C:
791... def __init__(self):
792... yield 1
793... def f(self):
794... yield 2
795>>> type(f())
796<type 'None'>
Tim Peters08a898f2001-06-28 01:52:22 +0000797
798>>> def f():
799... if 0:
800... return
801... if 0:
802... yield 2
803>>> type(f())
804<type 'generator'>
805
806
807>>> def f():
808... if 0:
809... lambda x: x # shouldn't trigger here
810... return # or here
811... def f(i):
812... return 2*i # or here
813... if 0:
814... return 3 # but *this* sucks (line 8)
815... if 0:
816... yield 2 # because it's a generator
817Traceback (most recent call last):
818SyntaxError: 'return' with argument inside generator (<string>, line 8)
Tim Petersea2e97a2001-06-24 07:10:02 +0000819"""
820
Tim Petersbe4f0a72001-06-29 02:41:16 +0000821# conjoin is a simple backtracking generator, named in honor of Icon's
822# "conjunction" control structure. Pass a list of no-argument functions
823# that return iterable objects. Easiest to explain by example: assume the
824# function list [x, y, z] is passed. Then conjoin acts like:
825#
826# def g():
827# values = [None] * 3
828# for values[0] in x():
829# for values[1] in y():
830# for values[2] in z():
831# yield values
832#
833# So some 3-lists of values *may* be generated, each time we successfully
834# get into the innermost loop. If an iterator fails (is exhausted) before
835# then, it "backtracks" to get the next value from the nearest enclosing
836# iterator (the one "to the left"), and starts all over again at the next
837# slot (pumps a fresh iterator). Of course this is most useful when the
838# iterators have side-effects, so that which values *can* be generated at
839# each slot depend on the values iterated at previous slots.
840
841def conjoin(gs):
842
843 values = [None] * len(gs)
844
845 def gen(i, values=values):
846 if i >= len(gs):
847 yield values
848 else:
849 for values[i] in gs[i]():
850 for x in gen(i+1):
851 yield x
852
853 for x in gen(0):
854 yield x
855
Tim Petersc468fd22001-06-30 07:29:44 +0000856# That works fine, but recursing a level and checking i against len(gs) for
857# each item produced is inefficient. By doing manual loop unrolling across
858# generator boundaries, it's possible to eliminate most of that overhead.
859# This isn't worth the bother *in general* for generators, but conjoin() is
860# a core building block for some CPU-intensive generator applications.
861
862def conjoin(gs):
863
864 n = len(gs)
865 values = [None] * n
866
867 # Do one loop nest at time recursively, until the # of loop nests
868 # remaining is divisible by 3.
869
870 def gen(i, values=values):
871 if i >= n:
872 yield values
873
874 elif (n-i) % 3:
875 ip1 = i+1
876 for values[i] in gs[i]():
877 for x in gen(ip1):
878 yield x
879
880 else:
881 for x in _gen3(i):
882 yield x
883
884 # Do three loop nests at a time, recursing only if at least three more
885 # remain. Don't call directly: this is an internal optimization for
886 # gen's use.
887
888 def _gen3(i, values=values):
889 assert i < n and (n-i) % 3 == 0
890 ip1, ip2, ip3 = i+1, i+2, i+3
891 g, g1, g2 = gs[i : ip3]
892
893 if ip3 >= n:
894 # These are the last three, so we can yield values directly.
895 for values[i] in g():
896 for values[ip1] in g1():
897 for values[ip2] in g2():
898 yield values
899
900 else:
901 # At least 6 loop nests remain; peel off 3 and recurse for the
902 # rest.
903 for values[i] in g():
904 for values[ip1] in g1():
905 for values[ip2] in g2():
906 for x in _gen3(ip3):
907 yield x
908
909 for x in gen(0):
910 yield x
911
Tim Petersbe4f0a72001-06-29 02:41:16 +0000912# A conjoin-based N-Queens solver.
913
914class Queens:
915 def __init__(self, n):
916 self.n = n
917 rangen = range(n)
918
919 # Assign a unique int to each column and diagonal.
920 # columns: n of those, range(n).
921 # NW-SE diagonals: 2n-1 of these, i-j unique and invariant along
922 # each, smallest i-j is 0-(n-1) = 1-n, so add n-1 to shift to 0-
923 # based.
924 # NE-SW diagonals: 2n-1 of these, i+j unique and invariant along
925 # each, smallest i+j is 0, largest is 2n-2.
926
927 # For each square, compute a bit vector of the columns and
928 # diagonals it covers, and for each row compute a function that
929 # generates the possiblities for the columns in that row.
930 self.rowgenerators = []
931 for i in rangen:
932 rowuses = [(1L << j) | # column ordinal
933 (1L << (n + i-j + n-1)) | # NW-SE ordinal
934 (1L << (n + 2*n-1 + i+j)) # NE-SW ordinal
935 for j in rangen]
936
937 def rowgen(rowuses=rowuses):
938 for j in rangen:
939 uses = rowuses[j]
Tim Petersc468fd22001-06-30 07:29:44 +0000940 if uses & self.used == 0:
941 self.used |= uses
942 yield j
943 self.used &= ~uses
Tim Petersbe4f0a72001-06-29 02:41:16 +0000944
945 self.rowgenerators.append(rowgen)
946
947 # Generate solutions.
948 def solve(self):
949 self.used = 0
950 for row2col in conjoin(self.rowgenerators):
951 yield row2col
952
953 def printsolution(self, row2col):
954 n = self.n
955 assert n == len(row2col)
956 sep = "+" + "-+" * n
957 print sep
958 for i in range(n):
959 squares = [" " for j in range(n)]
960 squares[row2col[i]] = "Q"
961 print "|" + "|".join(squares) + "|"
962 print sep
963
964conjoin_tests = """
965
966Generate the 3-bit binary numbers in order. This illustrates dumbest-
967possible use of conjoin, just to generate the full cross-product.
968
Tim Petersc468fd22001-06-30 07:29:44 +0000969>>> for c in conjoin([lambda: (0, 1)] * 3):
Tim Petersbe4f0a72001-06-29 02:41:16 +0000970... print c
971[0, 0, 0]
972[0, 0, 1]
973[0, 1, 0]
974[0, 1, 1]
975[1, 0, 0]
976[1, 0, 1]
977[1, 1, 0]
978[1, 1, 1]
979
Tim Petersc468fd22001-06-30 07:29:44 +0000980For efficiency in typical backtracking apps, conjoin() yields the same list
981object each time. So if you want to save away a full account of its
982generated sequence, you need to copy its results.
983
984>>> def gencopy(iterator):
985... for x in iterator:
986... yield x[:]
987
988>>> for n in range(10):
989... all = list(gencopy(conjoin([lambda: (0, 1)] * n)))
990... print n, len(all), all[0] == [0] * n, all[-1] == [1] * n
9910 1 1 1
9921 2 1 1
9932 4 1 1
9943 8 1 1
9954 16 1 1
9965 32 1 1
9976 64 1 1
9987 128 1 1
9998 256 1 1
10009 512 1 1
1001
Tim Petersbe4f0a72001-06-29 02:41:16 +00001002And run an 8-queens solver.
1003
1004>>> q = Queens(8)
1005>>> LIMIT = 2
1006>>> count = 0
1007>>> for row2col in q.solve():
1008... count += 1
1009... if count <= LIMIT:
1010... print "Solution", count
1011... q.printsolution(row2col)
1012Solution 1
1013+-+-+-+-+-+-+-+-+
1014|Q| | | | | | | |
1015+-+-+-+-+-+-+-+-+
1016| | | | |Q| | | |
1017+-+-+-+-+-+-+-+-+
1018| | | | | | | |Q|
1019+-+-+-+-+-+-+-+-+
1020| | | | | |Q| | |
1021+-+-+-+-+-+-+-+-+
1022| | |Q| | | | | |
1023+-+-+-+-+-+-+-+-+
1024| | | | | | |Q| |
1025+-+-+-+-+-+-+-+-+
1026| |Q| | | | | | |
1027+-+-+-+-+-+-+-+-+
1028| | | |Q| | | | |
1029+-+-+-+-+-+-+-+-+
1030Solution 2
1031+-+-+-+-+-+-+-+-+
1032|Q| | | | | | | |
1033+-+-+-+-+-+-+-+-+
1034| | | | | |Q| | |
1035+-+-+-+-+-+-+-+-+
1036| | | | | | | |Q|
1037+-+-+-+-+-+-+-+-+
1038| | |Q| | | | | |
1039+-+-+-+-+-+-+-+-+
1040| | | | | | |Q| |
1041+-+-+-+-+-+-+-+-+
1042| | | |Q| | | | |
1043+-+-+-+-+-+-+-+-+
1044| |Q| | | | | | |
1045+-+-+-+-+-+-+-+-+
1046| | | | |Q| | | |
1047+-+-+-+-+-+-+-+-+
1048
1049>>> print count, "solutions in all."
105092 solutions in all.
1051"""
1052
Tim Petersf6ed0742001-06-27 07:17:57 +00001053__test__ = {"tut": tutorial_tests,
1054 "pep": pep_tests,
1055 "email": email_tests,
1056 "fun": fun_tests,
Tim Petersbe4f0a72001-06-29 02:41:16 +00001057 "syntax": syntax_tests,
1058 "conjoin": conjoin_tests}
Tim Peters1def3512001-06-23 20:27:04 +00001059
1060# Magic test name that regrtest.py invokes *after* importing this module.
1061# This worms around a bootstrap problem.
1062# Note that doctest and regrtest both look in sys.argv for a "-v" argument,
1063# so this works as expected in both ways of running regrtest.
1064def test_main():
1065 import doctest, test_generators
Tim Peters2106ef02001-06-25 01:30:12 +00001066 if 0:
1067 # Temporary block to help track down leaks. So far, the blame
Tim Petersf6ed0742001-06-27 07:17:57 +00001068 # fell mostly on doctest. Later: the only leaks remaining are
1069 # in fun_tests, and only if you comment out the two LazyList.clear()
1070 # calls.
1071 for i in range(10000):
Tim Peters2106ef02001-06-25 01:30:12 +00001072 doctest.master = None
1073 doctest.testmod(test_generators)
1074 else:
1075 doctest.testmod(test_generators)
Tim Peters1def3512001-06-23 20:27:04 +00001076
1077# This part isn't needed for regrtest, but for running the test directly.
1078if __name__ == "__main__":
1079 test_main()