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