def test_join_on_key_phased(self): a = Chain(range(4)).map(lambda x: {'a': x, 'b': 'z%d' % x}) b = Chain(range(4)).map(lambda x: {'a': 3 - x, 'c': 'y%d' % x}) a.join_on_key('a', b) desired_result = [ { 'a': 0, 'b': 'z0', 'c': 'y3' }, { 'a': 1, 'b': 'z1', 'c': 'y2' }, { 'a': 2, 'b': 'z2', 'c': 'y1' }, { 'a': 3, 'b': 'z3', 'c': 'y0' }, ] actual_result = sorted(list(a), key=lambda o: o['a']) self.assertEqual(actual_result, desired_result)
def test_join_on_key(self): a = Chain(range(4)).map(lambda x: {'a': x, 'b': 'z%d' % x}) b = Chain(range(4)).map(lambda x: {'a': x, 'c': 'y%d' % x}) a.join_on_key('a', b) desired_result = [ { 'a': 0, 'b': 'z0', 'c': 'y0' }, { 'a': 1, 'b': 'z1', 'c': 'y1' }, { 'a': 2, 'b': 'z2', 'c': 'y2' }, { 'a': 3, 'b': 'z3', 'c': 'y3' }, ] self.assertEqual(list(a), desired_result)
def test_mapcat(self): def gen(x): for letter in ['a', 'b']: yield '%s%s' % (x, letter) b = Chain(range(3)).mapcat(gen) self.assertEqual(list(b), ['0a', '0b', '1a', '1b', '2a', '2b'])
def test_copy_after_map(self): a = Chain(range(3)).map(lambda x: 2 * x) b = a.copy() a.map(lambda x: 3 * x) self.assertEqual(list(a), [0, 6, 12]) self.assertEqual(list(b), [0, 2, 4])
def test_map_kwargs(self): def mult(x, m=2): return x * m a = range(3) b = Chain(a).map(mult, m=3) self.assertEqual(list(b), [0, 3, 6])
def test_copy(self): a = Chain(range(3)) b = a.copy() a.map(lambda x: 2 * x) self.assertEqual(list(a), [0, 2, 4]) self.assertEqual(list(b), [0, 1, 2])
def test_iter(self): a = Chain(range(2)) it = iter(a) self.assertEqual(next(it), 0) self.assertEqual(next(it), 1) with self.assertRaises(StopIteration): next(it)
def test_keep_keys_missing(self): objs = [ dict(a=1, b=2), dict(a=4, b=5), ] res = list(Chain(objs).keep_keys(['a', 'c'])) desired_result = [{'a': 1}, {'a': 4}] self.assertEqual(res, desired_result)
def test_for_each(self): res = [] def f(x): res.append(x) Chain(range(3)).for_each(f) self.assertEqual(res, [0, 1, 2])
def test_do_kwargs(self): res = [] def f(x, m): res.append(x * m) list(Chain(range(3)).do(f, m=3)) self.assertEqual(res, [0, 3, 6])
def test_do(self): res = [] def f(x): res.append(x) list(Chain(range(3)).do(f)) self.assertEqual(res, [0, 1, 2])
def test_omit_kwargs(self): a = range(5) def f(x, m): return x % m == 0 b = Chain(a).omit(f, m=3) self.assertEqual(list(b), [1, 2, 4])
def test_sink(self): res = [] def f(x): res.append(x) Chain(range(3)).do(f).sink() self.assertEqual(res, [0, 1, 2])
def test_keep_keys_multiple(self): objs = [ dict(a=1, b=2, c=3), dict(a=4, b=5, c=6), ] res = list(Chain(objs).keep_keys(['a', 'b'])) desired_result = [{'a': 1, 'b': 2}, {'a': 4, 'b': 5}] self.assertEqual(res, desired_result)
def test_transform_kwargs(self): def t(in_iterator, m): for x in in_iterator: if x % m == 0: yield x a = Chain(range(7)).transform(t, m=3) self.assertEqual(list(a), [0, 3, 6])
def test_transform(self): def t(in_iterator): for x in in_iterator: if x % 2 == 0: yield x a = Chain(range(5)).transform(t) self.assertEqual(list(a), [0, 2, 4])
def test_map_error_no_error_fn(self): def f(x): if x == 1: raise Exception('bad') return 2 * x b = Chain(range(4)).map(f) with self.assertRaises(Exception): list(b)
def test_set_key_fn_kwargs(self): objs = [{'a': 4}, {'a': 2}, {'a': 1}] def f(x, m): return x['a'] + m res = list(Chain(objs).set_key('b', f, m=3)) desired_result = [{'a': 4, 'b': 7}, {'a': 2, 'b': 5}, {'a': 1, 'b': 4}] self.assertEqual(res, desired_result)
def test_next(self): a = Chain(range(2)) self.assertEqual(next(a), 0) self.assertEqual(next(a), 1) try: next(a) except StopIteration: pass else: self.fail()
def test_flatten_not_strict(self): def err_f(ex, o): self.assertFalse('Should not call this!') def f(x): if x == 2: return x return range(x) b = Chain(range(4)).on_error(err_f).map(f).flatten(strict=False) self.assertEqual(list(b), [0, 2, 0, 1, 2])
def test_counter(self): l = [] def accumulate(count): l.append(count) Chain(range(4))\ .transform(utils.counter(accumulate))\ .filter(lambda x: x % 2 == 0)\ .transform(utils.counter(accumulate))\ .sink() self.assertEqual(l, [4, 2])
def test_flatten_strict(self): def err_f(ex, o): self.assertIsInstance(ex, TypeError) self.assertEqual(o, 2) def f(x): if x == 2: return x return range(x) b = Chain(range(4)).on_error(err_f).map(f).flatten() self.assertEqual(list(b), [0, 0, 1, 2])
def test_reduce_error(self): e = Exception('bad') def err_f(ex, o): self.assertEqual(ex, e) self.assertEqual(o, 1) def f(r, x): if x == 1: raise e return r + x res = Chain(range(4)).on_error(err_f).reduce(f) self.assertEqual(res, 5)
def test_do_error(self): e = Exception('bad') res = [] def err_f(ex, o): self.assertEqual(ex, e) self.assertEqual(o, 1) def f(x): if x == 1: raise e res.append(x) list(Chain(range(3)).do(f).on_error(err_f)) self.assertEqual(res, [0, 2])
def test_for_each_error(self): res = [] e = Exception('bad') def err_f(ex, o): self.assertEqual(ex, e) self.assertEqual(o, 1) def f(x): if x == 1: raise e res.append(x) Chain(range(3)).on_error(err_f).for_each(f) self.assertEqual(res, [0, 2])
def test_map_key(self): def make_obj(x): return dict(k=x) b = Chain(range(3)).map(make_obj).map_key('k', lambda x: x * 2) self.assertEqual(list(b), [ { 'k': 0 }, { 'k': 2 }, { 'k': 4 }, ])
def test_filter_error(self): a = range(4) e = Exception('bad') def err_f(ex, o): self.assertEqual(ex, e) self.assertEqual(o, 1) def f(x): if x == 1: raise e return x % 2 == 1 b = Chain(a)\ .on_error(err_f)\ .filter(f) self.assertEqual(list(b), [3])
def test_map_late_error(self): a = range(4) e = Exception('bad') def handle_error(ex, o): self.assertEqual(ex, e) self.assertEqual(o, 1) def f(x): if x == 1: raise e return 2 * x b = Chain(a)\ .map(f)\ .on_error(handle_error) self.assertEqual(list(b), [0, 4, 6])
def test_mapcat_error_2(self): e = Exception('bad') def err_f(ex, o): self.assertEqual(ex, e) self.assertEqual(o, 1) def gen(x): for letter in ['a', 'b']: out = '%s%s' % (x, letter) if out == '1b': raise e else: yield out b = Chain(range(3)).mapcat(gen).on_error(err_f) self.assertEqual(list(b), ['0a', '0b', '1a', '2a', '2b'])
def test_map_key_list(self): def make_obj(x): return dict(k=x, kk=x + 1) b = Chain(range(3)).map(make_obj).map_key(['k', 'kk'], lambda x: x * 2) self.assertEqual(list(b), [ { 'k': 0, 'kk': 2 }, { 'k': 2, 'kk': 4 }, { 'k': 4, 'kk': 6 }, ])