-
Notifications
You must be signed in to change notification settings - Fork 0
/
iterutil.py
78 lines (65 loc) · 2.47 KB
/
iterutil.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Utilities for iterables.
Created on Wed Nov 1 14:44:36 2017
@author: Juha Jeronen <juha.jeronen@tut.fi>
"""
# generator version
def uniqify(iterable, key=None):
"""Remove duplicates from iterable; preserve ordering.
If a key is provided, key(elt) is tested instead of elt for whether
a given element has already been seen.
"""
key = key or (lambda x: x)
it = iter(iterable)
seen = set()
for e in it:
k = key(e)
if k not in seen:
seen.add(k)
yield e
def flatten(iterable):
"""Remove nested structure from iterable.
http://rightfootin.blogspot.fi/2006/09/more-on-python-flatten.html
This version by Danny Yoo.
"""
it = iter(iterable)
for e in it:
if isinstance(e, (list, tuple)):
for f in flatten(e):
yield f
else:
yield e
def flatten1(iterable):
"""Remove nested structure from iterable, outermost level only."""
it = iter(iterable)
for e in it:
if isinstance(e, (list, tuple)):
for f in e:
yield f
else:
yield e
def flatten_if(iterable, condition):
"""Remove nested structure from iterable, matching items only.
E.g. to flatten only tuples of tuples (leaving tuples of atoms intact),
use something like:
condition = lambda item: isinstance(item[0], (tuple, list))
"""
it = iter(iterable)
for e in it:
if isinstance(e, (list, tuple)) and condition(e):
for f in flatten_if(e, condition):
yield f
else:
yield e
def test():
from operator import itemgetter
assert tuple(uniqify((1, 1, 2, 2, 2, 2, 4, 3, 3, 3))) == (1, 2, 4, 3)
assert tuple(uniqify((('foo', 1), ('bar', 1), ('foo', 2), ('baz', 2), ('qux', 4), ('foo', 3)), key=itemgetter(0))) == (('foo', 1), ('bar', 1), ('baz', 2), ('qux', 4))
assert tuple(uniqify((('foo', 1), ('bar', 1), ('foo', 2), ('baz', 2), ('qux', 4), ('foo', 3)), key=itemgetter(1))) == (('foo', 1), ('foo', 2), ('qux', 4), ('foo', 3))
assert tuple(flatten(((1, 2), (3, (4, 5), 6), (7, 8, 9)))) == (1, 2, 3, 4, 5, 6, 7, 8, 9)
assert tuple(flatten1(((1, 2), (3, (4, 5), 6), (7, 8, 9)))) == (1, 2, 3, (4, 5), 6, 7, 8, 9)
assert tuple(flatten_if((((1, 2), (3, 4)), (5, 6)), lambda e: isinstance(e[0], (tuple, list)))) == ((1, 2), (3, 4), (5, 6))
print("All tests passed")
if __name__ == '__main__':
test()