Learning Python - Mark Lutz [390]
>>> class stepper:
... def __getitem__(self, i):
... return self.data[i]
...
>>> X = stepper() # X is a stepper object
>>> X.data = "Spam"
>>>
>>> X[1] # Indexing calls __getitem__
'p'
>>> for item in X: # for loops call __getitem__
... print(item, end=' ') # for indexes items 0..N
...
S p a m
In fact, it’s really a case of “buy one, get a bunch free.” Any class that supports for loops automatically supports all iteration contexts in Python, many of which we’ve seen in earlier chapters (iteration contexts were presented in Chapter 14). For example, the in membership test, list comprehensions, the map built-in, list and tuple assignments, and type constructors will also call __getitem__ automatically, if it’s defined:
>>> 'p' in X # All call __getitem__ too
True
>>> [c for c in X] # List comprehension
['S', 'p', 'a', 'm']
>>> list(map(str.upper, X)) # map calls (use list() in 3.0)
['S', 'P', 'A', 'M']
>>> (a, b, c, d) = X # Sequence assignments
>>> a, c, d
('S', 'a', 'm')
>>> list(X), tuple(X), ''.join(X)
(['S', 'p', 'a', 'm'], ('S', 'p', 'a', 'm'), 'Spam')
>>> X
<__main__.stepper object at 0x00A8D5D0>
In practice, this technique can be used to create objects that provide a sequence interface and to add logic to built-in sequence type operations; we’ll revisit this idea when extending built-in types in Chapter 31.
Iterator Objects: __iter__ and __next__
Although the __getitem__ technique of the prior section works, it’s really just a fallback for iteration. Today, all iteration contexts in Python will try the __iter__ method first, before trying __getitem__. That is, they prefer the iteration protocol we learned about in Chapter 14 to repeatedly indexing an object; only if the object does not support the iteration protocol is indexing attempted instead. Generally speaking, you should prefer __iter__ too—it supports general iteration contexts better than __getitem__ can.
Technically, iteration contexts work by calling the iter built-in function to try to find an __iter__ method, which is expected to return an iterator object. If it’s provided, Python then repeatedly calls this iterator object’s __next__ method to produce items until a StopIteration exception is raised. If no such __iter__ method is found, Python falls back on the __getitem__ scheme and repeatedly indexes by offsets as before, until an IndexError exception is raised. A next built-in function is also available as a convenience for manual iterations: next(I) is the same as I.__next__().
* * *
Note
Version skew note: As described in Chapter 14, if you are using Python 2.6, the I.__next__() method just described is named I.next() in your Python, and the next(I) built-in is present for portability: it calls I.next() in 2.6 and I.__next__() in 3.0. Iteration works the same in 2.6 in all other respects.
* * *
User-Defined Iterators
In the __iter__ scheme, classes implement user-defined iterators by simply implementing the iteration protocol introduced in Chapters 14 and 20 (refer back to those chapters for more background details on iterators). For example, the following file, iters.py, defines a user-defined iterator class that generates squares:
class Squares:
def __init__(self, start, stop): # Save state when created
self.value = start - 1
self.stop = stop
def __iter__(self): # Get iterator object on iter
return self
def __next__(self): # Return a square on each iteration
if self.value == self.stop: # Also called by next built-in
raise StopIteration
self.value += 1
return self.value ** 2
% python
>>> from iters import Squares
>>> for i in Squares(1, 5): # for calls iter, which calls __iter__
... print(i, end=' ') # Each iteration calls __next__
...
1 4 9 16 25
Here, the iterator object is simply the instance self, because the __next__ method is part of this class. In more complex scenarios, the iterator