ansible-later/testenv/lib/python2.7/site-packages/toolz/tests/test_itertoolz.py
2019-04-23 13:04:27 +02:00

527 lines
17 KiB
Python

import itertools
from itertools import starmap
from toolz.utils import raises
from functools import partial
from random import Random
from pickle import dumps, loads
from toolz.itertoolz import (remove, groupby, merge_sorted,
concat, concatv, interleave, unique,
isiterable, getter,
mapcat, isdistinct, first, second,
nth, take, tail, drop, interpose, get,
rest, last, cons, frequencies,
reduceby, iterate, accumulate,
sliding_window, count, partition,
partition_all, take_nth, pluck, join,
diff, topk, peek, random_sample)
from toolz.compatibility import range, filter
from operator import add, mul
# is comparison will fail between this and no_default
no_default2 = loads(dumps('__no__default__'))
def identity(x):
return x
def iseven(x):
return x % 2 == 0
def isodd(x):
return x % 2 == 1
def inc(x):
return x + 1
def double(x):
return 2 * x
def test_remove():
r = remove(iseven, range(5))
assert type(r) is not list
assert list(r) == list(filter(isodd, range(5)))
def test_groupby():
assert groupby(iseven, [1, 2, 3, 4]) == {True: [2, 4], False: [1, 3]}
def test_groupby_non_callable():
assert groupby(0, [(1, 2), (1, 3), (2, 2), (2, 4)]) == \
{1: [(1, 2), (1, 3)],
2: [(2, 2), (2, 4)]}
assert groupby([0], [(1, 2), (1, 3), (2, 2), (2, 4)]) == \
{(1,): [(1, 2), (1, 3)],
(2,): [(2, 2), (2, 4)]}
assert groupby([0, 0], [(1, 2), (1, 3), (2, 2), (2, 4)]) == \
{(1, 1): [(1, 2), (1, 3)],
(2, 2): [(2, 2), (2, 4)]}
def test_merge_sorted():
assert list(merge_sorted([1, 2, 3], [1, 2, 3])) == [1, 1, 2, 2, 3, 3]
assert list(merge_sorted([1, 3, 5], [2, 4, 6])) == [1, 2, 3, 4, 5, 6]
assert list(merge_sorted([1], [2, 4], [3], [])) == [1, 2, 3, 4]
assert list(merge_sorted([5, 3, 1], [6, 4, 3], [],
key=lambda x: -x)) == [6, 5, 4, 3, 3, 1]
assert list(merge_sorted([2, 1, 3], [1, 2, 3],
key=lambda x: x // 3)) == [2, 1, 1, 2, 3, 3]
assert list(merge_sorted([2, 3], [1, 3],
key=lambda x: x // 3)) == [2, 1, 3, 3]
assert ''.join(merge_sorted('abc', 'abc', 'abc')) == 'aaabbbccc'
assert ''.join(merge_sorted('abc', 'abc', 'abc', key=ord)) == 'aaabbbccc'
assert ''.join(merge_sorted('cba', 'cba', 'cba',
key=lambda x: -ord(x))) == 'cccbbbaaa'
assert list(merge_sorted([1], [2, 3, 4], key=identity)) == [1, 2, 3, 4]
data = [[(1, 2), (0, 4), (3, 6)], [(5, 3), (6, 5), (8, 8)],
[(9, 1), (9, 8), (9, 9)]]
assert list(merge_sorted(*data, key=lambda x: x[1])) == [
(9, 1), (1, 2), (5, 3), (0, 4), (6, 5), (3, 6), (8, 8), (9, 8), (9, 9)]
assert list(merge_sorted()) == []
assert list(merge_sorted([1, 2, 3])) == [1, 2, 3]
assert list(merge_sorted([1, 4, 5], [2, 3])) == [1, 2, 3, 4, 5]
assert list(merge_sorted([1, 4, 5], [2, 3], key=identity)) == [
1, 2, 3, 4, 5]
assert list(merge_sorted([1, 5], [2], [4, 7], [3, 6], key=identity)) == [
1, 2, 3, 4, 5, 6, 7]
def test_interleave():
assert ''.join(interleave(('ABC', '123'))) == 'A1B2C3'
assert ''.join(interleave(('ABC', '1'))) == 'A1BC'
def test_unique():
assert tuple(unique((1, 2, 3))) == (1, 2, 3)
assert tuple(unique((1, 2, 1, 3))) == (1, 2, 3)
assert tuple(unique((1, 2, 3), key=iseven)) == (1, 2)
def test_isiterable():
assert isiterable([1, 2, 3]) is True
assert isiterable('abc') is True
assert isiterable(5) is False
def test_isdistinct():
assert isdistinct([1, 2, 3]) is True
assert isdistinct([1, 2, 1]) is False
assert isdistinct("Hello") is False
assert isdistinct("World") is True
assert isdistinct(iter([1, 2, 3])) is True
assert isdistinct(iter([1, 2, 1])) is False
def test_nth():
assert nth(2, 'ABCDE') == 'C'
assert nth(2, iter('ABCDE')) == 'C'
assert nth(1, (3, 2, 1)) == 2
assert nth(0, {'foo': 'bar'}) == 'foo'
assert raises(StopIteration, lambda: nth(10, {10: 'foo'}))
assert nth(-2, 'ABCDE') == 'D'
assert raises(ValueError, lambda: nth(-2, iter('ABCDE')))
def test_first():
assert first('ABCDE') == 'A'
assert first((3, 2, 1)) == 3
assert isinstance(first({0: 'zero', 1: 'one'}), int)
def test_second():
assert second('ABCDE') == 'B'
assert second((3, 2, 1)) == 2
assert isinstance(second({0: 'zero', 1: 'one'}), int)
def test_last():
assert last('ABCDE') == 'E'
assert last((3, 2, 1)) == 1
assert isinstance(last({0: 'zero', 1: 'one'}), int)
def test_rest():
assert list(rest('ABCDE')) == list('BCDE')
assert list(rest((3, 2, 1))) == list((2, 1))
def test_take():
assert list(take(3, 'ABCDE')) == list('ABC')
assert list(take(2, (3, 2, 1))) == list((3, 2))
def test_tail():
assert list(tail(3, 'ABCDE')) == list('CDE')
assert list(tail(3, iter('ABCDE'))) == list('CDE')
assert list(tail(2, (3, 2, 1))) == list((2, 1))
def test_drop():
assert list(drop(3, 'ABCDE')) == list('DE')
assert list(drop(1, (3, 2, 1))) == list((2, 1))
def test_take_nth():
assert list(take_nth(2, 'ABCDE')) == list('ACE')
def test_get():
assert get(1, 'ABCDE') == 'B'
assert list(get([1, 3], 'ABCDE')) == list('BD')
assert get('a', {'a': 1, 'b': 2, 'c': 3}) == 1
assert get(['a', 'b'], {'a': 1, 'b': 2, 'c': 3}) == (1, 2)
assert get('foo', {}, default='bar') == 'bar'
assert get({}, [1, 2, 3], default='bar') == 'bar'
assert get([0, 2], 'AB', 'C') == ('A', 'C')
assert get([0], 'AB') == ('A',)
assert get([], 'AB') == ()
assert raises(IndexError, lambda: get(10, 'ABC'))
assert raises(KeyError, lambda: get(10, {'a': 1}))
assert raises(TypeError, lambda: get({}, [1, 2, 3]))
assert raises(TypeError, lambda: get([1, 2, 3], 1, None))
assert raises(KeyError, lambda: get('foo', {}, default=no_default2))
def test_mapcat():
assert (list(mapcat(identity, [[1, 2, 3], [4, 5, 6]])) ==
[1, 2, 3, 4, 5, 6])
assert (list(mapcat(reversed, [[3, 2, 1, 0], [6, 5, 4], [9, 8, 7]])) ==
list(range(10)))
inc = lambda i: i + 1
assert ([4, 5, 6, 7, 8, 9] ==
list(mapcat(partial(map, inc), [[3, 4, 5], [6, 7, 8]])))
def test_cons():
assert list(cons(1, [2, 3])) == [1, 2, 3]
def test_concat():
assert list(concat([[], [], []])) == []
assert (list(take(5, concat([['a', 'b'], range(1000000000)]))) ==
['a', 'b', 0, 1, 2])
def test_concatv():
assert list(concatv([], [], [])) == []
assert (list(take(5, concatv(['a', 'b'], range(1000000000)))) ==
['a', 'b', 0, 1, 2])
def test_interpose():
assert "a" == first(rest(interpose("a", range(1000000000))))
assert "tXaXrXzXaXn" == "".join(interpose("X", "tarzan"))
assert list(interpose(0, itertools.repeat(1, 4))) == [1, 0, 1, 0, 1, 0, 1]
assert list(interpose('.', ['a', 'b', 'c'])) == ['a', '.', 'b', '.', 'c']
def test_frequencies():
assert (frequencies(["cat", "pig", "cat", "eel",
"pig", "dog", "dog", "dog"]) ==
{"cat": 2, "eel": 1, "pig": 2, "dog": 3})
assert frequencies([]) == {}
assert frequencies("onomatopoeia") == {"a": 2, "e": 1, "i": 1, "m": 1,
"o": 4, "n": 1, "p": 1, "t": 1}
def test_reduceby():
data = [1, 2, 3, 4, 5]
iseven = lambda x: x % 2 == 0
assert reduceby(iseven, add, data, 0) == {False: 9, True: 6}
assert reduceby(iseven, mul, data, 1) == {False: 15, True: 8}
projects = [{'name': 'build roads', 'state': 'CA', 'cost': 1000000},
{'name': 'fight crime', 'state': 'IL', 'cost': 100000},
{'name': 'help farmers', 'state': 'IL', 'cost': 2000000},
{'name': 'help farmers', 'state': 'CA', 'cost': 200000}]
assert reduceby(lambda x: x['state'],
lambda acc, x: acc + x['cost'],
projects, 0) == {'CA': 1200000, 'IL': 2100000}
assert reduceby('state',
lambda acc, x: acc + x['cost'],
projects, 0) == {'CA': 1200000, 'IL': 2100000}
def test_reduce_by_init():
assert reduceby(iseven, add, [1, 2, 3, 4]) == {True: 2 + 4, False: 1 + 3}
assert reduceby(iseven, add, [1, 2, 3, 4], no_default2) == {True: 2 + 4,
False: 1 + 3}
def test_reduce_by_callable_default():
def set_add(s, i):
s.add(i)
return s
assert reduceby(iseven, set_add, [1, 2, 3, 4, 1, 2], set) == \
{True: set([2, 4]), False: set([1, 3])}
def test_iterate():
assert list(itertools.islice(iterate(inc, 0), 0, 5)) == [0, 1, 2, 3, 4]
assert list(take(4, iterate(double, 1))) == [1, 2, 4, 8]
def test_accumulate():
assert list(accumulate(add, [1, 2, 3, 4, 5])) == [1, 3, 6, 10, 15]
assert list(accumulate(mul, [1, 2, 3, 4, 5])) == [1, 2, 6, 24, 120]
assert list(accumulate(add, [1, 2, 3, 4, 5], -1)) == [-1, 0, 2, 5, 9, 14]
def binop(a, b):
raise AssertionError('binop should not be called')
start = object()
assert list(accumulate(binop, [], start)) == [start]
assert list(accumulate(add, [1, 2, 3], no_default2)) == [1, 3, 6]
def test_accumulate_works_on_consumable_iterables():
assert list(accumulate(add, iter((1, 2, 3)))) == [1, 3, 6]
def test_sliding_window():
assert list(sliding_window(2, [1, 2, 3, 4])) == [(1, 2), (2, 3), (3, 4)]
assert list(sliding_window(3, [1, 2, 3, 4])) == [(1, 2, 3), (2, 3, 4)]
def test_sliding_window_of_short_iterator():
assert list(sliding_window(3, [1, 2])) == []
def test_partition():
assert list(partition(2, [1, 2, 3, 4])) == [(1, 2), (3, 4)]
assert list(partition(3, range(7))) == [(0, 1, 2), (3, 4, 5)]
assert list(partition(3, range(4), pad=-1)) == [(0, 1, 2),
(3, -1, -1)]
assert list(partition(2, [])) == []
def test_partition_all():
assert list(partition_all(2, [1, 2, 3, 4])) == [(1, 2), (3, 4)]
assert list(partition_all(3, range(5))) == [(0, 1, 2), (3, 4)]
assert list(partition_all(2, [])) == []
def test_count():
assert count((1, 2, 3)) == 3
assert count([]) == 0
assert count(iter((1, 2, 3, 4))) == 4
assert count('hello') == 5
assert count(iter('hello')) == 5
def test_pluck():
assert list(pluck(0, [[0, 1], [2, 3], [4, 5]])) == [0, 2, 4]
assert list(pluck([0, 1], [[0, 1, 2], [3, 4, 5]])) == [(0, 1), (3, 4)]
assert list(pluck(1, [[0], [0, 1]], None)) == [None, 1]
data = [{'id': 1, 'name': 'cheese'}, {'id': 2, 'name': 'pies', 'price': 1}]
assert list(pluck('id', data)) == [1, 2]
assert list(pluck('price', data, 0)) == [0, 1]
assert list(pluck(['id', 'name'], data)) == [(1, 'cheese'), (2, 'pies')]
assert list(pluck(['name'], data)) == [('cheese',), ('pies',)]
assert list(pluck(['price', 'other'], data, 0)) == [(0, 0), (1, 0)]
assert raises(IndexError, lambda: list(pluck(1, [[0]])))
assert raises(KeyError, lambda: list(pluck('name', [{'id': 1}])))
assert list(pluck(0, [[0, 1], [2, 3], [4, 5]], no_default2)) == [0, 2, 4]
assert raises(IndexError, lambda: list(pluck(1, [[0]], no_default2)))
def test_join():
names = [(1, 'one'), (2, 'two'), (3, 'three')]
fruit = [('apple', 1), ('orange', 1), ('banana', 2), ('coconut', 2)]
def addpair(pair):
return pair[0] + pair[1]
result = set(starmap(add, join(first, names, second, fruit)))
expected = set([((1, 'one', 'apple', 1)),
((1, 'one', 'orange', 1)),
((2, 'two', 'banana', 2)),
((2, 'two', 'coconut', 2))])
assert result == expected
result = set(starmap(add, join(first, names, second, fruit,
left_default=no_default2,
right_default=no_default2)))
assert result == expected
def test_getter():
assert getter(0)('Alice') == 'A'
assert getter([0])('Alice') == ('A',)
assert getter([])('Alice') == ()
def test_key_as_getter():
squares = [(i, i**2) for i in range(5)]
pows = [(i, i**2, i**3) for i in range(5)]
assert set(join(0, squares, 0, pows)) == set(join(lambda x: x[0], squares,
lambda x: x[0], pows))
get = lambda x: (x[0], x[1])
assert set(join([0, 1], squares, [0, 1], pows)) == set(join(get, squares,
get, pows))
get = lambda x: (x[0],)
assert set(join([0], squares, [0], pows)) == set(join(get, squares,
get, pows))
def test_join_double_repeats():
names = [(1, 'one'), (2, 'two'), (3, 'three'), (1, 'uno'), (2, 'dos')]
fruit = [('apple', 1), ('orange', 1), ('banana', 2), ('coconut', 2)]
result = set(starmap(add, join(first, names, second, fruit)))
expected = set([((1, 'one', 'apple', 1)),
((1, 'one', 'orange', 1)),
((2, 'two', 'banana', 2)),
((2, 'two', 'coconut', 2)),
((1, 'uno', 'apple', 1)),
((1, 'uno', 'orange', 1)),
((2, 'dos', 'banana', 2)),
((2, 'dos', 'coconut', 2))])
assert result == expected
def test_join_missing_element():
names = [(1, 'one'), (2, 'two'), (3, 'three')]
fruit = [('apple', 5), ('orange', 1)]
result = set(starmap(add, join(first, names, second, fruit)))
expected = set([((1, 'one', 'orange', 1))])
assert result == expected
def test_left_outer_join():
result = set(join(identity, [1, 2], identity, [2, 3], left_default=None))
expected = set([(2, 2), (None, 3)])
assert result == expected
def test_right_outer_join():
result = set(join(identity, [1, 2], identity, [2, 3], right_default=None))
expected = set([(2, 2), (1, None)])
assert result == expected
def test_outer_join():
result = set(join(identity, [1, 2], identity, [2, 3],
left_default=None, right_default=None))
expected = set([(2, 2), (1, None), (None, 3)])
assert result == expected
def test_diff():
assert raises(TypeError, lambda: list(diff()))
assert raises(TypeError, lambda: list(diff([1, 2])))
assert raises(TypeError, lambda: list(diff([1, 2], 3)))
assert list(diff([1, 2], (1, 2), iter([1, 2]))) == []
assert list(diff([1, 2, 3], (1, 10, 3), iter([1, 2, 10]))) == [
(2, 10, 2), (3, 3, 10)]
assert list(diff([1, 2], [10])) == [(1, 10)]
assert list(diff([1, 2], [10], default=None)) == [(1, 10), (2, None)]
# non-variadic usage
assert raises(TypeError, lambda: list(diff([])))
assert raises(TypeError, lambda: list(diff([[]])))
assert raises(TypeError, lambda: list(diff([[1, 2]])))
assert raises(TypeError, lambda: list(diff([[1, 2], 3])))
assert list(diff([(1, 2), (1, 3)])) == [(2, 3)]
data1 = [{'cost': 1, 'currency': 'dollar'},
{'cost': 2, 'currency': 'dollar'}]
data2 = [{'cost': 100, 'currency': 'yen'},
{'cost': 300, 'currency': 'yen'}]
conversions = {'dollar': 1, 'yen': 0.01}
def indollars(item):
return conversions[item['currency']] * item['cost']
list(diff(data1, data2, key=indollars)) == [
({'cost': 2, 'currency': 'dollar'}, {'cost': 300, 'currency': 'yen'})]
def test_topk():
assert topk(2, [4, 1, 5, 2]) == (5, 4)
assert topk(2, [4, 1, 5, 2], key=lambda x: -x) == (1, 2)
assert topk(2, iter([5, 1, 4, 2]), key=lambda x: -x) == (1, 2)
assert topk(2, [{'a': 1, 'b': 10}, {'a': 2, 'b': 9},
{'a': 10, 'b': 1}, {'a': 9, 'b': 2}], key='a') == \
({'a': 10, 'b': 1}, {'a': 9, 'b': 2})
assert topk(2, [{'a': 1, 'b': 10}, {'a': 2, 'b': 9},
{'a': 10, 'b': 1}, {'a': 9, 'b': 2}], key='b') == \
({'a': 1, 'b': 10}, {'a': 2, 'b': 9})
assert topk(2, [(0, 4), (1, 3), (2, 2), (3, 1), (4, 0)], 0) == \
((4, 0), (3, 1))
def test_topk_is_stable():
assert topk(4, [5, 9, 2, 1, 5, 3], key=lambda x: 1) == (5, 9, 2, 1)
def test_peek():
alist = ["Alice", "Bob", "Carol"]
element, blist = peek(alist)
element == alist[0]
assert list(blist) == alist
assert raises(StopIteration, lambda: peek([]))
def test_random_sample():
alist = list(range(100))
assert list(random_sample(prob=1, seq=alist, random_state=2016)) == alist
mk_rsample = lambda rs=1: list(random_sample(prob=0.1,
seq=alist,
random_state=rs))
rsample1 = mk_rsample()
assert rsample1 == mk_rsample()
rsample2 = mk_rsample(1984)
randobj = Random(1984)
assert rsample2 == mk_rsample(randobj)
assert rsample1 != rsample2
assert mk_rsample(object) == mk_rsample(object)
assert mk_rsample(object) != mk_rsample(object())
assert mk_rsample(b"a") == mk_rsample(u"a")
assert raises(TypeError, lambda: mk_rsample([]))