diff --git a/swh/core/tests/test_utils.py b/swh/core/tests/test_utils.py new file mode 100644 index 0000000..1475eaa --- /dev/null +++ b/swh/core/tests/test_utils.py @@ -0,0 +1,25 @@ +# Copyright (C) 2015-2016 The Software Heritage developers +# See the AUTHORS file at the top-level directory of this distribution +# License: GNU General Public License version 3, or any later version +# See top-level LICENSE file for more information + +import unittest + +from nose.tools import istest + +from swh.core import utils + + +class UtilsTest(unittest.TestCase): + @istest + def grouper(self): + # given + gen = (d for d in [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]) + + # when + actual_group = [] + for data in utils.grouper(gen, 3): + actual_group.append(list(data)) + + # then + self.assertEquals(actual_group, [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]]) diff --git a/swh/core/utils.py b/swh/core/utils.py new file mode 100644 index 0000000..0559d96 --- /dev/null +++ b/swh/core/utils.py @@ -0,0 +1,26 @@ +# Copyright (C) 2015-2016 The Software Heritage developers +# See the AUTHORS file at the top-level directory of this distribution +# License: GNU General Public License version 3, or any later version +# See top-level LICENSE file for more information + + +import itertools + + +def grouper(iterable, n): + """Collect data into fixed-length chunks or blocks. + The last block is exactly the size of the remaining data. + + Args: + iterable: an iterable + n: size of block + + Returns: + fixed-length chunks of blocks as iterables (except for the last + one which can be of size < n) + + """ + args = [iter(iterable)] * n + fv = None + for _data in itertools.zip_longest(*args, fillvalue=fv): + yield (d for d in _data if d is not fv)