# coding=utf-8

"""
itertools模块
"""

import itertools
import unittest

class TestItertools(unittest.TestCase):

    def test_count(self):
        """itertools.count()无限产生自然数"""
        natuals = itertools.count(1, 1)
        res = 0
        for r in natuals:
            if r == 10:
                res = r
                break
        self.assertEqual(10, res)

    def test_cycle(self):
        """itertools.cycle()无限重复序列里面的值"""
        my_list = "ABC"
        result = []
        for num, res in enumerate(itertools.cycle(my_list)):
            result.append(res)
            if num == 5:
                break
        self.assertEqual([‘A‘,‘B‘,‘C‘,‘A‘,‘B‘,‘C‘], result)

    def test_repeat(self):
        """itertools.repeat(‘A‘,3) 重复A 3次"""
        result = [res for res in itertools.repeat(‘ABC‘, 3)]
        self.assertEqual([‘ABC‘, ‘ABC‘, ‘ABC‘], result)

    def test_takewhile(self):
        """itertools.takewhile()截取一个有限的序列"""
        natuals = itertools.count(0)
        result = itertools.takewhile(lambda x : x <= 5, natuals)
        self.assertEqual([0, 1, 2, 3, 4, 5], list(result))

    def test_chain(self):
        """itertools.chain()将两个迭代对象组成一个迭代对象"""
        iter1 = [‘A‘,‘B‘]
        iter2 = ‘CCC‘
        result = itertools.chain(iter1, iter2)
        self.assertEqual([‘A‘, ‘B‘, ‘C‘, ‘C‘, ‘C‘], list(result))

    def test_groupby(self):
        """groupby()把迭代器中相邻的重复元素挑出来放在一起"""
        text = "AAAAAACCCCDDDDMMM"
        result = {key : list(val) for key, val in itertools.groupby(text)}
        self.assertEqual([‘M‘, ‘M‘, ‘M‘], result[‘M‘])

if __name__ == ‘__main__‘:
    unittest.main()