Lib/test/test_itertools.py

   1 import unittest
   2 from test import test_support
   3 from itertools import *
   4 import sys
   5
   6 class TestBasicOps(unittest.TestCase):
   7     def test_chain(self):
   8         self.assertEqual(list(chain('abc', 'def')), list('abcdef'))
   9
  10     def test_count(self):
  11         self.assertEqual(zip('abc',count()), [('a', 0), ('b', 1), ('c', 2)])
  12         self.assertEqual(zip('abc',count(3)), [('a', 3), ('b', 4), ('c', 5)])
  13         self.assertRaises(TypeError, count, 2, 3)
  14
  15     def test_cycle(self):
  16         self.assertEqual(list(islice(cycle('abc'),10)), list('abcabcabca'))
  17         self.assertEqual(list(cycle('')), [])
  18
  19     def test_ifilter(self):
  20         def isEven(x):
  21             return x%2==0
  22         self.assertEqual(list(ifilter(isEven, range(6))), [0,2,4])
  23         self.assertEqual(list(ifilter(None, [0,1,0,2,0])), [1,2])
  24         self.assertRaises(TypeError, ifilter)
  25         self.assertRaises(TypeError, ifilter, 3)
  26         self.assertRaises(TypeError, ifilter, isEven, 3)
  27
  28     def test_ifilterfalse(self):
  29         def isEven(x):
  30             return x%2==0
  31         self.assertEqual(list(ifilterfalse(isEven, range(6))), [1,3,5])
  32         self.assertEqual(list(ifilterfalse(None, [0,1,0,2,0])), [0,0,0])
  33         self.assertRaises(TypeError, ifilterfalse)
  34         self.assertRaises(TypeError, ifilterfalse, 3)
  35         self.assertRaises(TypeError, ifilterfalse, isEven, 3)
  36
  37     def test_izip(self):
  38         ans = [(x,y) for x, y in izip('abc',count())]
  39         self.assertEqual(ans, [('a', 0), ('b', 1), ('c', 2)])
  40         self.assertRaises(TypeError, izip)
  41
  42     def test_repeat(self):
  43         self.assertEqual(zip(xrange(3),repeat('a')),
  44                          [(0, 'a'), (1, 'a'), (2, 'a')])
  45         self.assertEqual(list(repeat('a', 3)), ['a', 'a', 'a'])
  46         self.assertRaises(TypeError, repeat)
  47
  48     def test_imap(self):
  49         import operator
  50         self.assertEqual(list(imap(operator.pow, range(3), range(1,7))),
  51                          [0**1, 1**2, 2**3])
  52         self.assertEqual(list(imap(None, 'abc', range(5))),
  53                          [('a',0),('b',1),('c',2)])
  54         self.assertRaises(TypeError, imap)
  55         self.assertRaises(TypeError, imap, operator.neg)
  56
  57     def test_starmap(self):
  58         import operator
  59         self.assertEqual(list(starmap(operator.pow, zip(range(3), range(1,7)))),
  60                          [0**1, 1**2, 2**3])
  61         self.assertRaises(TypeError, list, starmap(operator.pow, [[4,5]]))
  62
  63     def test_islice(self):
  64         for args in [          # islice(args) should agree with range(args)
  65                 (10, 20, 3),
  66                 (10, 3, 20),
  67                 (10, 20),
  68                 (10, 3),
  69                 (20,)
  70                 ]:
  71             self.assertEqual(list(islice(xrange(100), *args)), range(*args))
  72
  73         for args, tgtargs in [  # Stop when seqn is exhausted
  74                 ((10, 110, 3), ((10, 100, 3))),
  75                 ((10, 110), ((10, 100))),
  76                 ((110,), (100,))
  77                 ]:
  78             self.assertEqual(list(islice(xrange(100), *args)), range(*tgtargs))
  79
  80         self.assertRaises(TypeError, islice, xrange(10))
  81         self.assertRaises(TypeError, islice, xrange(10), 1, 2, 3, 4)
  82         self.assertRaises(ValueError, islice, xrange(10), -5, 10, 1)
  83         self.assertRaises(ValueError, islice, xrange(10), 1, -5, -1)
  84         self.assertRaises(ValueError, islice, xrange(10), 1, 10, -1)
  85         self.assertRaises(ValueError, islice, xrange(10), 1, 10, 0)
  86         self.assertEqual(len(list(islice(count(), 1, 10, sys.maxint))), 1)
  87
  88     def test_takewhile(self):
  89         data = [1, 3, 5, 20, 2, 4, 6, 8]
  90         underten = lambda x: x<10
  91         self.assertEqual(list(takewhile(underten, data)), [1, 3, 5])
  92
  93     def test_dropwhile(self):
  94         data = [1, 3, 5, 20, 2, 4, 6, 8]
  95         underten = lambda x: x<10
  96         self.assertEqual(list(dropwhile(underten, data)), [20, 2, 4, 6, 8])
  97
  98 libreftest = """ Doctest for examples in the library reference, libitertools.tex
  99
 100
 101 >>> amounts = [120.15, 764.05, 823.14]
 102 >>> for checknum, amount in izip(count(1200), amounts):
 103 ...     print 'Check %d is for $%.2f' % (checknum, amount)
 104 ...
 105 Check 1200 is for $120.15
 106 Check 1201 is for $764.05
 107 Check 1202 is for $823.14
 108
 109 >>> import operator
 110 >>> import operator
 111 >>> for cube in imap(operator.pow, xrange(1,4), repeat(3)):
 112 ...    print cube
 113 ...
 114 1
 115 8
 116 27
 117
 118 >>> reportlines = ['EuroPython', 'Roster', '', 'alex', '', 'laura', '', 'martin', '', 'walter', '', 'samuele']
 119 >>> for name in islice(reportlines, 3, len(reportlines), 2):
 120 ...    print name.title()
 121 ...
 122 Alex
 123 Laura
 124 Martin
 125 Walter
 126 Samuele
 127
 128 >>> def enumerate(iterable):
 129 ...     return izip(count(), iterable)
 130
 131 >>> def tabulate(function):
 132 ...     "Return function(0), function(1), ..."
 133 ...     return imap(function, count())
 134
 135 >>> def iteritems(mapping):
 136 ...     return izip(mapping.iterkeys(), mapping.itervalues())
 137
 138 >>> def nth(iterable, n):
 139 ...     "Returns the nth item"
 140 ...     return list(islice(iterable, n, n+1))
 141
 142 >>> def all(pred, seq):
 143 ...     "Returns True if pred(x) is True for every element in the iterable"
 144 ...     return not nth(ifilterfalse(pred, seq), 0)
 145
 146 >>> def some(pred, seq):
 147 ...     "Returns True if pred(x) is True at least one element in the iterable"
 148 ...     return bool(nth(ifilter(pred, seq), 0))
 149
 150 >>> def no(pred, seq):
 151 ...     "Returns True if pred(x) is False for every element in the iterable"
 152 ...     return not nth(ifilter(pred, seq), 0)
 153
 154 >>> def pairwise(seq):
 155 ...     "s -> (s0,s1), (s1,s2), (s2, s3), ..."
 156 ...     return izip(seq, islice(seq,1,len(seq)))
 157
 158 """
 159
 160 __test__ = {'libreftest' : libreftest}
 161
 162 def test_main(verbose=None):
 163     import test_itertools
 164     suite = unittest.TestSuite()
 165     for testclass in (TestBasicOps,
 166                           ):
 167         suite.addTest(unittest.makeSuite(testclass))
 168     test_support.run_suite(suite)
 169     test_support.run_doctest(test_itertools, verbose)
 170
 171     # verify reference counting
 172     import sys
 173     if verbose and hasattr(sys, "gettotalrefcount"):
 174         counts = []
 175         for i in xrange(5):
 176             test_support.run_suite(suite)
 177             counts.append(sys.gettotalrefcount()-i)
 178         print counts
 179
 180 if __name__ == "__main__":
 181     test_main(verbose=True)