python/mozbuild/mozbuild/chunkify.py

   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
  10 from itertools import islice
  11
  12
  13 class ChunkingError(Exception):
  14     pass
  15
  16
  17 def 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
  45 def 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)