In python, I am trying to write a class that support two different kind of iterator. Roughly speaking, this object contains a matrix of data and I want to have two different kind of iterator to support row iteration and column iteration.
Okay, so make two separate methods where each is a generator.
class Matrix(object):
def iter_rows(self):
for row in self.rows:
yield row
def iter_columns(self):
for column in self.columns:
yield column
Your __iter__
could iterate over one or the other by default, though I'd recommend having no __iter__
at all.
Is this what you're looking for?
class Matrix(object):
def __init__(self, rows):
self._rows = rows
def columns(self):
return zip(*self._rows)
def rows(self):
return self._rows
# Create a Matrix by providing rows.
m = Matrix([[1,2,3],
[4,5,6],
[7,8,9]])
# Iterate in row-major order.
for row in m.rows():
for value in row:
print value
# Iterate in column-major order.
for column in m.columns():
for value in column:
print value
You can use itertools.izip
instead of zip
if you want to create each column on demand.
You can also move the iteration of the actual values into the class. I wasn't sure if you wanted to iterate over rows/columns (as shown) or the values in the rows/columns.
dict
has several iterator-producing methods -- iterkeys
, itervalues
, iteritems
-- and so should your class. If there's one "most natural" way of iterating, you should also alias it to __iter__
for convenience and readability (that's probably going to be iterrows
; of course there is always going to be some doubt, as there was with dict
when we designed its iteration behavior, but a reasonable pick is better than none).
For example, suppose your matrix is square, held flattened up into a row-major list self.data
, with a side of self.n
. Then:
def iterrows(self):
start = 0
n = self.n
data = self.data
while start < n*n:
stop = start + n
yield data[start:stop]
start = stop
def itercols(self):
start = 0
n = self.n
data = self.data
while start < n:
yield data[start::n]
start += 1
__iter__ = iterrows