Add doctests to some of this stuff.
[python_utils.git] / list_utils.py
1 #!/usr/bin/env python3
2
3 from itertools import islice
4 from typing import Any, Iterator, List
5
6
7 def shard(lst: List[Any], size: int) -> Iterator[Any]:
8     """
9     Yield successive size-sized shards from lst.
10
11     >>> for sublist in shard([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12], 3):
12     ...     [_ for _ in sublist]
13     [1, 2, 3]
14     [4, 5, 6]
15     [7, 8, 9]
16     [10, 11, 12]
17
18     """
19     for x in range(0, len(lst), size):
20         yield islice(lst, x, x + size)
21
22
23 def flatten(lst: List[Any]) -> List[Any]:
24     """
25     Flatten out a list:
26
27     >>> flatten([ 1, [2, 3, 4, [5], 6], 7, [8, [9]]])
28     [1, 2, 3, 4, 5, 6, 7, 8, 9]
29
30     """
31     if len(lst) == 0:
32         return lst
33     if isinstance(lst[0], list):
34         return flatten(lst[0]) + flatten(lst[1:])
35     return lst[:1] + flatten(lst[1:])
36
37
38 def prepend(item: Any, lst: List[Any]) -> List[Any]:
39     """
40     Prepend an item to a list.
41
42     >>> prepend('foo', ['bar', 'baz'])
43     ['foo', 'bar', 'baz']
44
45     """
46     lst.insert(0, item)
47     return lst
48
49
50 if __name__ == '__main__':
51     import doctest
52     doctest.testmod()