diff --git a/toolz/itertoolz.py b/toolz/itertoolz.py index 354ecf25..c5e37343 100644 --- a/toolz/itertoolz.py +++ b/toolz/itertoolz.py @@ -319,6 +319,8 @@ def take(n, seq): drop tail """ + if n < 0: + raise ValueError('take: n must be a non-negative integer, got %r' % (n,)) return itertools.islice(seq, n) @@ -348,6 +350,8 @@ def drop(n, seq): take tail """ + if n < 0: + raise ValueError('drop: n must be a non-negative integer, got %r' % (n,)) return itertools.islice(seq, n, None) diff --git a/toolz/tests/test_itertoolz.py b/toolz/tests/test_itertoolz.py index c8640917..f756d89f 100644 --- a/toolz/tests/test_itertoolz.py +++ b/toolz/tests/test_itertoolz.py @@ -185,6 +185,15 @@ def test_take(): assert list(take(2, (3, 2, 1))) == list((3, 2)) +def test_take_negative_n(): + try: + list(take(-1, [1, 2, 3])) + assert False, 'expected ValueError' + except ValueError as e: + assert 'non-negative' in str(e) + assert list(take(0, [1, 2, 3])) == [] # n == 0 boundary still returns, not rejected + + def test_tail(): assert list(tail(3, 'ABCDE')) == list('CDE') assert list(tail(3, iter('ABCDE'))) == list('CDE') @@ -196,6 +205,15 @@ def test_drop(): assert list(drop(1, (3, 2, 1))) == list((2, 1)) +def test_drop_negative_n(): + try: + list(drop(-1, [1, 2, 3])) + assert False, 'expected ValueError' + except ValueError as e: + assert 'non-negative' in str(e) + assert list(drop(0, [1, 2, 3])) == [1, 2, 3] + + def test_take_nth(): assert list(take_nth(2, 'ABCDE')) == list('ACE')