How would one create an iterative function (or iterator object) in python?
Iterator objects in python conform to the iterator protocol, which basically means they provide two methods: __iter__()
and __next__()
.
The __iter__
returns the iterator object and is implicitly called
at the start of loops.
The __next__()
method returns the next value and is implicitly called at each loop increment. This method raises a StopIteration exception when there are no more value to return, which is implicitly captured by looping constructs to stop iterating.
Here's a simple example of a counter:
class Counter:
def __init__(self, low, high):
self.current = low - 1
self.high = high
def __iter__(self):
return self
def __next__(self): # Python 2: def next(self)
self.current += 1
if self.current < self.high:
return self.current
raise StopIteration
for c in Counter(3, 9):
print(c)
This will print:
3
4
5
6
7
8
This is easier to write using a generator, as covered in a previous answer:
def counter(low, high):
current = low
while current < high:
yield current
current += 1
for c in counter(3, 9):
print(c)
The printed output will be the same. Under the hood, the generator object supports the iterator protocol and does something roughly similar to the class Counter.
David Mertz's article, Iterators and Simple Generators, is a pretty good introduction.
Inspired by Matt Gregory's answer here is a bit more complicated iterator that will return a,b,...,z,aa,ab,...,zz,aaa,aab,...,zzy,zzz
class AlphaCounter:
def __init__(self, low, high):
self.current = low
self.high = high
def __iter__(self):
return self
def __next__(self): # Python 3: def __next__(self)
alpha = ' abcdefghijklmnopqrstuvwxyz'
n_current = sum([(alpha.find(self.current[x])* 26**(len(self.current)-x-1)) for x in range(len(self.current))])
n_high = sum([(alpha.find(self.high[x])* 26**(len(self.high)-x-1)) for x in range(len(self.high))])
if n_current > n_high:
raise StopIteration
else:
increment = True
ret = ''
for x in self.current[::-1]:
if 'z' == x:
if increment:
ret += 'a'
else:
ret += 'z'
else:
if increment:
ret += alpha[alpha.find(x)+1]
increment = False
else:
ret += x
if increment:
ret += 'a'
tmp = self.current
self.current = ret[::-1]
return tmp
for c in AlphaCounter('a', 'zzz'):
print(c)
This question is about iterable objects, not about iterators. In Python, sequences are iterable too so one way to make an iterable class is to make it behave like a sequence, i.e. give it __getitem__
and __len__
methods. I have tested this on Python 2 and 3.
class CustomRange:
def __init__(self, low, high):
self.low = low
self.high = high
def __getitem__(self, item):
if item >= len(self):
raise IndexError("CustomRange index out of range")
return self.low + item
def __len__(self):
return self.high - self.low
cr = CustomRange(0, 10)
for i in cr:
print(i)
First of all the itertools module is incredibly useful for all sorts of cases in which an iterator would be useful, but here is all you need to create an iterator in python:
yield
Isn't that cool? Yield can be used to replace a normal return in a function. It returns the object just the same, but instead of destroying state and exiting, it saves state for when you want to execute the next iteration. Here is an example of it in action pulled directly from the itertools function list:
def count(n=0):
while True:
yield n
n += 1
As stated in the functions description (it's the count() function from the itertools module...) , it produces an iterator that returns consecutive integers starting with n.
Generator expressions are a whole other can of worms (awesome worms!). They may be used in place of a List Comprehension to save memory (list comprehensions create a list in memory that is destroyed after use if not assigned to a variable, but generator expressions can create a Generator Object... which is a fancy way of saying Iterator). Here is an example of a generator expression definition:
gen = (n for n in xrange(0,11))
This is very similar to our iterator definition above except the full range is predetermined to be between 0 and 10.
I just found xrange() (suprised I hadn't seen it before...) and added it to the above example. xrange() is an iterable version of range() which has the advantage of not prebuilding the list. It would be very useful if you had a giant corpus of data to iterate over and only had so much memory to do it in.
I see some of you doing return self
in __iter__
. I just wanted to note that __iter__
itself can be a generator (thus removing the need for __next__
and raising StopIteration
exceptions)
class range:
def __init__(self,a,b):
self.a = a
self.b = b
def __iter__(self):
i = self.a
while i < self.b:
yield i
i+=1
Of course here one might as well directly make a generator, but for more complex classes it can be useful.
All answers on this page are really great for a complex object. But for those containing builtin iterable types as attributes, like str
, list
, set
or dict
, or any implementation of collections.Iterable
, you can omit certain things in your class.
class Test(object):
def __init__(self, string):
self.string = string
def __iter__(self):
# since your string is already iterable
return (ch for ch in self.string)
# or simply
return self.string.__iter__()
# also
return iter(self.string)
It can be used like:
for x in Test("abcde"):
print(x)
# prints
# a
# b
# c
# d
# e