| 936 | |
| 937 | class TestExtend(unittest.TestCase): |
| 938 | def test_custom_functions(self): |
| 939 | @extend(aslist=True) |
| 940 | def my_zip(it): |
| 941 | return zip(it, it) |
| 942 | |
| 943 | result = seq.range(3).my_zip().list() |
| 944 | expected = list(zip(range(3), range(3))) |
| 945 | self.assertEqual(result, expected) |
| 946 | |
| 947 | result = seq.range(3).my_zip().my_zip().list() |
| 948 | expected = list(zip(expected, expected)) |
| 949 | self.assertEqual(result, expected) |
| 950 | |
| 951 | @extend |
| 952 | def square(it): |
| 953 | return [i ** 2 for i in it] |
| 954 | |
| 955 | result = seq.range(100).square().list() |
| 956 | expected = [i ** 2 for i in range(100)] |
| 957 | self.assertEqual(result, expected) |
| 958 | |
| 959 | name = "PARALLEL_SQUARE" |
| 960 | |
| 961 | @extend(parallel=True, name=name) |
| 962 | def square_parallel(it): |
| 963 | return [i ** 2 for i in it] |
| 964 | |
| 965 | result = seq.range(100).square_parallel() |
| 966 | self.assertEqual(result.sum(), sum(expected)) |
| 967 | self.assertEqual( |
| 968 | repr(result._lineage), "Lineage: sequence -> extended[%s]" % name |
| 969 | ) |
| 970 | |
| 971 | @extend |
| 972 | def my_filter(it, n=10): |
| 973 | return (i for i in it if i > n) |
| 974 | |
| 975 | # test keyword args |
| 976 | result = seq.range(20).my_filter(n=10).list() |
| 977 | expected = list(filter(lambda x: x > 10, range(20))) |
| 978 | self.assertEqual(result, expected) |
| 979 | |
| 980 | # test args |
| 981 | result = seq.range(20).my_filter(10).list() |
| 982 | self.assertEqual(result, expected) |
| 983 | |
| 984 | # test final |
| 985 | @extend(final=True) |
| 986 | def toarray(it): |
| 987 | return array.array("f", it) |
| 988 | |
| 989 | result = seq.range(10).toarray() |
| 990 | expected = array.array("f", range(10)) |
| 991 | self.assertEqual(result, expected) |
| 992 | |
| 993 | result = seq.range(10).map(lambda x: x ** 2).toarray() |
| 994 | expected = array.array("f", [i ** 2 for i in range(10)]) |
| 995 | self.assertEqual(result, expected) |