"""
Extracts specified columns from a list by taking a list of column indices and converting them in to mininum number of Python
slice objects ahead of time before iterating through the list. Then uses slices to extract ranges of columns.
"""
from itertools import chain
def prepare_slices(indices):
"""
Converts a list of Python indices into an optimized list of slice objects.
>>> x = [2, 4, 13, 16, 19, 23, 24, 25]
>>> prepare_slices(x)
[slice(2, 5, 2), slice(13, 20, 3), slice(23, 26, 1)]
"""
start = None
try:
start = indices.pop(0)
stop = indices.pop(0)
except IndexError:
return [slice(start, None)]
step = stop - start
slices = []
for index in indices:
if step:
stride = index - stop
if stride == step:
stop = index
else:
slices.append(slice(start, stop + 1, step))
start, stop, step = index, None, None
else:
stop, step = index, index - start
slices.append(slice(start, stop + 1 if stop else None, step))
return slices
def slice_columns(seq, slices):
"""
Extracts items from a sequence using a list of slice objects
>>> letters_a_to_z = [chr(97 + i) for i in range(26)]
>>> slices_ = [slice(2, 5, 2), slice(13, 20, 3), slice(23, 26, 1)]
>>> slice_columns(letters_a_to_z, slices_)
['c', 'e', 'n', 'q', 't', 'x', 'y', 'z']
"""
return list(chain.from_iterable((seq[i] for i in slices)))
def slice_rows(rows, indices):
"""
Extracts columns from a list of rows
"""
slices = prepare_slices(indices)
for row in rows:
yield slice_columns(row, slices)
if __name__ == '__main__':
import doctest
doctest.testmod()