forum/utils/lists.py

   1 """Utilities for working with lists and sequences."""
   2
   3 def flatten(x):
   4     """
   5     Returns a single, flat list which contains all elements retrieved
   6     from the sequence and all recursively contained sub-sequences
   7     (iterables).
   8
   9     Examples:
  10     >>> [1, 2, [3, 4], (5, 6)]
  11     [1, 2, [3, 4], (5, 6)]
  12
  13     From http://kogs-www.informatik.uni-hamburg.de/~meine/python_tricks
  14     """
  15     result = []
  16     for el in x:
  17         if hasattr(el, '__iter__') and not isinstance(el, basestring):
  18             result.extend(flatten(el))
  19         else:
  20             result.append(el)
  21     return result
  22
  23 def batch_size(items, size):
  24     """
  25     Retrieves items in batches of the given size.
  26
  27     >>> l = range(1, 11)
  28     >>> batch_size(l, 3)
  29     [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10]]
  30     >>> batch_size(l, 5)
  31     [[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]
  32     """
  33     return [items[i:i+size] for i in xrange(0, len(items), size)]
  34
  35 def batches(items, number):
  36     """
  37     Retrieves items in the given number of batches.
  38
  39     >>> l = range(1, 11)
  40     >>> batches(l, 1)
  41     [[1, 2, 3, 4, 5, 6, 7, 8, 9, 10]]
  42     >>> batches(l, 2)
  43     [[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]
  44     >>> batches(l, 3)
  45     [[1, 2, 3, 4], [5, 6, 7, 8], [9, 10]]
  46     >>> batches(l, 4)
  47     [[1, 2, 3], [4, 5, 6], [7, 8, 9], [10]]
  48     >>> batches(l, 5)
  49     [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]]
  50
  51     Initial batches will contain as many items as possible in cases where
  52     there are not enough items to be distributed evenly.
  53
  54     >>> batches(l, 6)
  55     [[1, 2], [3, 4], [5, 6], [7, 8], [9], [10]]
  56     >>> batches(l, 7)
  57     [[1, 2], [3, 4], [5, 6], [7], [8], [9], [10]]
  58     >>> batches(l, 8)
  59     [[1, 2], [3, 4], [5], [6], [7], [8], [9], [10]]
  60     >>> batches(l, 9)
  61     [[1, 2], [3], [4], [5], [6], [7], [8], [9], [10]]
  62     >>> batches(l, 10)
  63     [[1], [2], [3], [4], [5], [6], [7], [8], [9], [10]]
  64
  65     If there are more batches than items, empty batches will be appended
  66     to the batch list.
  67
  68     >>> batches(l, 11)
  69     [[1], [2], [3], [4], [5], [6], [7], [8], [9], [10], []]
  70     >>> batches(l, 12)
  71     [[1], [2], [3], [4], [5], [6], [7], [8], [9], [10], [], []]
  72     """
  73     div, mod= divmod(len(items), number)
  74     if div > 1:
  75         if mod:
  76             div += 1
  77         return batch_size(items, div)
  78     else:
  79         if not div:
  80             return [[item] for item in items] + [[]] * (number - mod)
  81         elif div == 1 and not mod:
  82             return [[item] for item in items]
  83         else:
  84             # mod now tells you how many lists of 2 you can fit in
  85             return ([items[i*2:(i*2)+2] for i in xrange(0, mod)] +
  86                     [[item] for item in items[mod*2:]])