We can improve readability of your code as follows:
using type annotations partially looks incomplete, when we can be more descriptive:
from typing import (Iterable,
TypeVar)
ElementType = TypeVar('ElementType')
def iter_group(iterable: Iterable[ElementType],
batch_size: int) -> Iterable[Iterable[ElementType]]:
...
checking if object has given type should be done using isinstance built-in function, so instead of
iterable_type = type(iterable)
...
if iterable_type == list:
...
we can simply write
if isinstance(iterable, list):
...
(btw I don't understand why are you treating list as a special case)
if you want negate a number then we can use - unary operator, no need to multiply by -1, so instead of
start = batch_size * -1
we can write
start = -batch_size
no need in parentheses for while-loop condition, we are not in C/C++/JAVA, we can simply write:
while end < length:
...
setting start to -batch_size and end to 0 and change them right after that seems redundant when we can set them to 0 and batch_size respectfully and increment at the end of the while loop body, so instead
start = batch_size * -1
end = 0
while end < length:
start += batch_size
end += batch_size
...
we can write
start = 0
end = batch_size
while end < length:
...
start += batch_size
end += batch_size
If I understood correctly what you are trying to achieve is to split iterable into evenly sized chunks, which is a well-known StackOverflow question, but accepted answer works only for sequences (str/list/tuple) and so does your solution. For example it won't work for potentially infinite sequences like classic Fibonacci numbers generator
>>> def fibonacci():
a, b = 0, 1
while True:
yield a
a, b = b, a + b
>>> next(iter_group(fibonacci(), 5))
Traceback (most recent call last):
...
length = len(iterable)
TypeError: object of type 'generator' has no len()
If you want to work with arbitrary iterables (e.g. generators), we may use itertools module and this brilliant solution which uses iter built-in function form with sentinel value
from itertools import islice
from typing import (Iterable,
Tuple,
TypeVar)
ElementType = TypeVar('ElementType')
def chunks(iterable: Iterable[ElementType],
batch_size: int) -> Iterable[Tuple[ElementType, ...]]:
iterator = iter(iterable)
return iter(lambda: tuple(islice(iterator, batch_size)), ())
but if we make a simple benchmark with
def sequence_chunks(iterable: Sequence[ElementType],
batch_size: int) -> Iterable[Sequence[ElementType]]:
for start in range(0, len(iterable), batch_size):
yield iterable[start:start + batch_size]
like
import timeit
...
print('original solution',
min(timeit.repeat('list(iter_group(iterable, 1000))',
'from __main__ import iter_group\n'
'iterable = range(0, 10001)',
number=10000)))
print('Ned Batchelder\'s solution',
min(timeit.repeat('list(sequence_chunks(iterable, 1000))',
'from __main__ import sequence_chunks\n'
'iterable = range(0, 10001)',
number=10000)))
print('senderle\'s solution',
min(timeit.repeat('list(chunks(iterable, 1000))',
'from __main__ import chunks\n'
'iterable = range(0, 10001)',
number=10000)))
on my laptop with Windows 10 and Python 3.5 we'll have
original solution 0.07320549999999999
Ned Batchelder's solution 0.06249870000000002
senderle's solution 2.6072023999999994
so how can we have both speed and handle cases with non-sequence iterables?
Here comes functools module with singledispatch function decorator. We can use it like
import timeit
from collections import abc
from functools import singledispatch
from itertools import islice
from typing import (Iterable,
Sequence,
TypeVar)
ElementType = TypeVar('ElementType')
@singledispatch
def chunks(iterable: Iterable[ElementType],
batch_size: int) -> Iterable[Iterable[ElementType]]:
iterator = iter(iterable)
return iter(lambda: tuple(islice(iterator, batch_size)), ())
@chunks.register(abc.Sequence)
def sequence_chunks(iterable: Sequence[ElementType],
batch_size: int) -> Iterable[Iterable[ElementType]]:
for start in range(0, len(iterable), batch_size):
yield iterable[start:start + batch_size]
after that call to chunks will end up in sequence_chunks for sequences and in general chunks for all other cases.
But
print('single-dispatched solution',
min(timeit.repeat('list(chunks(iterable, 1000))',
'from __main__ import chunks\n'
'iterable = range(0, 10001)',
number=10000)))
gives
single-dispatched solution 0.0737681
so we lose some time during dispatching, but we can save some time and space if we will use itertools.islice for sequences as well like
@chunks.register(abc.Sequence)
def sequence_chunks(iterable: Sequence[ElementType],
batch_size: int) -> Iterable[Iterable[ElementType]]:
iterator = iter(iterable)
for _ in range(ceil_division(len(iterable), batch_size)):
yield islice(iterator, batch_size)
def ceil_division(left_number: int, right_number: int) -> int:
"""
Divides given numbers with ceiling.
"""
# based on https://stackoverflow.com/a/17511341/5997596
return -(-left_number // right_number)
which gives
complete single-dispatched solution 0.03895900000000002
P.S.
@AJNeufeld's solution gives
0.0647120000000001
on my laptop
