more itertools defined (google#15)

* added compress, cycle, dropwhile, from_iterable, ifilter, ifilterfalse, and takewhile in itertools
* added itertools tests for cycle, dropwhile, from_iterable, ifilter, ifilterfalse, takewhile

.gitignore

@@ -1,3 +1,4 @@
 build
 errors.err
 *.swp
+*.pyc

lib/itertools.py

@@ -17,19 +17,56 @@
 import _collections
 import sys
 
-
 def chain(*iterables):
   for it in iterables:
     for element in it:
       yield element
 
+def compress(data, selectors):
+  return (d for d,s in izip(data, selectors) if s)
 
 def count(start=0, step=1):
   n = start
   while True:
     yield n
     n += step
 
+def cycle(iterable):
+  saved = []
+  for element in iterable:
+    yield element
+    saved.append(element)
+  while saved:
+    for element in saved:
+      yield element
+
+def dropwhile(predicate, iterable):
+  iterable = iter(iterable)
+  for x in iterable:
+    if not predicate(x):
+      yield x
+      break
+  for x in iterable:
+    yield x
+
+def from_iterable(iterables):
+  for it in iterables:
+    for element in it:
+      yield element
+
+def ifilter(predicate, iterable):
+  if predicate is None:
+    predicate = bool
+  for x in iterable:
+    if predicate(x):
+       yield x
+
+def ifilterfalse(predicate, iterable):
+  if predicate is None:
+    predicate = bool
+  for x in iterable:
+    if not predicate(x):
+       yield x
 
 def imap(function, *iterables):
   iterables = map(iter, iterables)
@@ -40,7 +77,6 @@ def imap(function, *iterables):
     else:
       yield function(*args)
 
-
 def islice(iterable, *args):
   s = slice(*args)
   it = iter(xrange(s.start or 0, s.stop or sys.maxint, s.step or 1))
@@ -50,13 +86,11 @@ def islice(iterable, *args):
       yield element
       nexti = next(it)
 
-
 def izip(*iterables):
   iterators = map(iter, iterables)
   while iterators:
     yield tuple(map(next, iterators))
 
-
 def repeat(object, times=None):
   if times is None:
     while True:
@@ -65,11 +99,16 @@ def repeat(object, times=None):
     for i in xrange(times):
       yield object
 
-
 def starmap(function, iterable):
   for args in iterable:
     yield function(*args)
 
+def takewhile(predicate, iterable):
+  for x in iterable:
+    if predicate(x):
+      yield x
+    else:
+      break
 
 def tee(iterable, n=2):
   it = iter(iterable)

lib/itertools_test.py

@@ -16,6 +16,70 @@
 
 import weetest
 
+def TestCycle():
+  want = []
+  got = []
+  for x in itertools.cycle(()):
+      got.append(x)
+  assert got == want, 'empty cycle yields no elements'
+
+  arg = (0, 1, 2)
+  want = (0, 1, 2) * 10
+  got = []
+  limit = 10 * len(arg)
+  counter = 0
+  for x in itertools.cycle((0, 1, 2)):
+    got.append(x)
+    counter += 1
+    if counter == limit:
+      break
+  assert tuple(got) == want, 'tuple(cycle%s) == %s, want %s' % (arg, tuple(got), want)
+
+def TestDropwhile():
+  r = range(10)
+  cases = [
+    ((lambda x: x < 5, r), (5, 6, 7, 8, 9)),
+    ((lambda x: True, r), ()),
+    ((lambda x: False, r), tuple(r)),
+  ]
+  for args, want in cases:
+    got = tuple(itertools.dropwhile(*args))
+    assert got == want, 'tuple(dropwhile%s) == %s, want %s' % (args, got, want)
+
+def TestFromIterable():
+  r = range(10)
+  cases = [
+    ([r], tuple(r)),
+    ([r, r], tuple(r) + tuple(r)),
+    ([], ())
+  ]
+  for args, want in cases:
+    got = tuple(itertools.from_iterable(args))
+    assert got == want, 'tuple(from_iterable%s) == %s, want %s' % (args, got, want)
+
+def TestIFilter():
+  r = range(10)
+  cases = [
+    ((lambda x: x < 5, r), (0, 1, 2, 3, 4)),
+    ((lambda x: False, r), ()),
+    ((lambda x: True, r), tuple(r)),
+    ((None, r), (1, 2, 3, 4, 5, 6, 7, 8, 9))
+  ]
+  for args, want in cases:
+    got = tuple(itertools.ifilter(*args))
+    assert got == want, 'tuple(ifilter%s) == %s, want %s' % (args, got, want)
+
+def TestIFilterFalse():
+  r = range(10)
+  cases = [
+    ((lambda x: x < 5, r), (5, 6, 7, 8, 9)),
+    ((lambda x: False, r), tuple(r)),
+    ((lambda x: True, r), ()),
+    ((None, r), (0,))
+  ]
+  for args, want in cases:
+    got = tuple(itertools.ifilterfalse(*args))
+    assert got == want, 'tuple(ifilterfalse%s) == %s, want %s' % (args, got, want)
 
 def TestISlice():
   r = range(10)
@@ -28,6 +92,16 @@ def TestISlice():
     got = tuple(itertools.islice(*args))
     assert got == want, 'tuple(islice%s) == %s, want %s' % (args, got, want)
 
+def TestTakewhile():
+  r = range(10)
+  cases = [
+    ((lambda x: x % 2 == 0, r), (0,)),
+    ((lambda x: True, r), tuple(r)),
+    ((lambda x: False, r), ())
+  ]
+  for args, want in cases:
+    got = tuple(itertools.takewhile(*args))
+    assert got == want, 'tuple(takewhile%s) == %s, want %s' % (args, got, want)
 
 if __name__ == '__main__':
   weetest.RunTests()