def test_batch_on_lists():
assert list(batches([1, 2, 3, 4, 5, 6], 1)) == [[1], [2], [3], [4], [5],
[6]]
assert list(batches([1, 2, 3, 4, 5, 6], 2)) == [[1, 2], [3, 4], [5, 6]]
assert list(batches([1, 2, 3, 4, 5, 6], 3)) == [[1, 2, 3], [4, 5, 6]]
assert list(batches([1, 2, 3, 4, 5, 6], 4)) == [
[1, 2, 3, 4],
[5, 6],
]
def test_batch_order():
iterable = range(100)
batch_size = 2
output = batches(iterable, batch_size)
assert list(chain.from_iterable(output)) == list(iterable)
def test_batch_sizes():
iterable = range(100)
batch_size = 2
output = list(batches(iterable, batch_size))
for batch in output[:-1]:
assert len(batch) == batch_size
assert len(output[-1]) <= batch_size
def test_batch_with_loop():
iterable = [1, 2, 3, 4, 5, 6]
samples = {
# even batches
1: [[1], [2], [3], [4], [5], [6]],
2: [[1, 2], [3, 4], [5, 6]],
3: [[1, 2, 3], [4, 5, 6]],
# batches with rest
4: [[1, 2, 3, 4], [5, 6]],
}
for batch_size, expected in samples.items():
assert list(batches(iterable, batch_size)) == expected
from batch import batches
from pprint import pprint
# 4 Samples of features
example_features = [['F11', 'F12', 'F13', 'F14'], ['F21', 'F22', 'F23', 'F24'],
['F31', 'F32', 'F33', 'F34'], ['F41', 'F42', 'F43', 'F44']]
# 4 Samples of labels
example_labels = [['L11', 'L12'], ['L21', 'L22'], ['L31', 'L32'],
['L41', 'L42']]
# PPrint prints data structures like 2d arrays, so they are easier to read
pprint(batches(2, example_features, example_labels))
def test_batch_parametrized(batch_size, expected):
iterable = [1, 2, 3, 4, 5, 6]
assert list(batches(iterable, batch_size)) == expected