Skip to content

Commit b4f0450

Browse files
authored
Merge pull request #463 from ml31415/flatten
Fix #462 flattening of iterators without len
2 parents b3812a6 + 0edfa40 commit b4f0450

File tree

2 files changed

+17
-9
lines changed

2 files changed

+17
-9
lines changed

streamz/core.py

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
import asyncio
22
from collections import deque, defaultdict
33
from datetime import timedelta
4+
from itertools import chain
45
import functools
56
import logging
67
import six
@@ -1637,15 +1638,20 @@ class flatten(Stream):
16371638
"""
16381639
def update(self, x, who=None, metadata=None):
16391640
L = []
1640-
for i, item in enumerate(x):
1641-
if i == len(x) - 1:
1642-
y = self._emit(item, metadata=metadata)
1643-
else:
1644-
y = self._emit(item)
1641+
items = chain(x)
1642+
item = next(items)
1643+
for item_next in items:
1644+
y = self._emit(item)
1645+
item = item_next
16451646
if type(y) is list:
16461647
L.extend(y)
16471648
else:
16481649
L.append(y)
1650+
y = self._emit(item, metadata=metadata)
1651+
if type(y) is list:
1652+
L.extend(y)
1653+
else:
1654+
L.append(y)
16491655
return L
16501656

16511657

streamz/tests/test_core.py

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -800,13 +800,15 @@ def test_frequencies():
800800
assert L[-1] == {'a': 2, 'b': 1}
801801

802802

803-
def test_flatten():
803+
@pytest.mark.parametrize("iterators",
804+
[[[1, 2, 3], [4, 5], [6, 7, 8]],
805+
[(i for i in range(1, 7)), (i for i in range(7, 9))]])
806+
def test_flatten(iterators):
804807
source = Stream()
805808
L = source.flatten().sink_to_list()
806809

807-
source.emit([1, 2, 3])
808-
source.emit([4, 5])
809-
source.emit([6, 7, 8])
810+
for iterator in iterators:
811+
source.emit(iterator)
810812

811813
assert L == [1, 2, 3, 4, 5, 6, 7, 8]
812814

0 commit comments

Comments
 (0)