diff --git a/judge/utils/iterator.py b/judge/utils/iterator.py new file mode 100644 index 0000000000..19ca285fd3 --- /dev/null +++ b/judge/utils/iterator.py @@ -0,0 +1,7 @@ +from itertools import zip_longest + + +def chunk(iterable, size): + fill = object() + for group in zip_longest(*[iter(iterable)] * size, fillvalue=fill): + yield [item for item in group if item is not fill] diff --git a/judge/utils/tests/test_iterator.py b/judge/utils/tests/test_iterator.py new file mode 100644 index 0000000000..30153616be --- /dev/null +++ b/judge/utils/tests/test_iterator.py @@ -0,0 +1,18 @@ +import unittest +from itertools import chain + +from judge.utils.iterator import chunk + + +class ChunkTestCase(unittest.TestCase): + def test_empty(self): + self.assertEqual(list(chunk([], 5)), []) + + def test_normal(self): + for size in [10, 13, 100, 200]: + with self.subTest(f'chunk size {size}'): + result = list(chunk(range(100), size)) + self.assertEqual(list(chain(*result)), list(range(100))) + for part in result[:-1]: + self.assertEqual(len(part), size) + self.assertLessEqual(len(result[-1]), size)