1# This Source Code Form is subject to the terms of the Mozilla Public
2# License, v. 2.0. If a copy of the MPL was not distributed with this file,
3# You can obtain one at http://mozilla.org/MPL/2.0/.
4
5# This file is a direct clone of
6# https://github.com/bhearsum/chunkify/blob/master/chunkify/__init__.py
7# of version 1.2. Its license (MPL2) is contained in repo root LICENSE file.
8# Please make modifications there where possible.
9
10from itertools import islice
11
12
13class ChunkingError(Exception):
14    pass
15
16
17def split_evenly(n, chunks):
18    """Split an integer into evenly distributed list
19
20    >>> split_evenly(7, 3)
21    [3, 2, 2]
22
23    >>> split_evenly(12, 3)
24    [4, 4, 4]
25
26    >>> split_evenly(35, 10)
27    [4, 4, 4, 4, 4, 3, 3, 3, 3, 3]
28
29    >>> split_evenly(1, 2)
30    Traceback (most recent call last):
31        ...
32    ChunkingError: Number of chunks is greater than number
33
34    """
35    if n < chunks:
36        raise ChunkingError("Number of chunks is greater than number")
37    if n % chunks == 0:
38        # Either we can evenly split or only 1 chunk left
39        return [n / chunks] * chunks
40    # otherwise the current chunk should be a bit larger
41    max_size = n / chunks + 1
42    return [max_size] + split_evenly(n - max_size, chunks - 1)
43
44
45def chunkify(things, this_chunk, chunks):
46    if this_chunk > chunks:
47        raise ChunkingError("this_chunk is greater than total chunks")
48
49    dist = split_evenly(len(things), chunks)
50    start = sum(dist[:this_chunk-1])
51    end = start + dist[this_chunk-1]
52
53    try:
54        return things[start:end]
55    except TypeError:
56        return islice(things, start, end)
57
58