91 lines
2.8 KiB
Python
91 lines
2.8 KiB
Python
from asyncio import Future
|
|
from threading import RLock
|
|
from typing import Any, List, Optional, Tuple
|
|
|
|
from reactivex import Observable, abc, from_future
|
|
from reactivex.disposable import CompositeDisposable, SingleAssignmentDisposable
|
|
from reactivex.internal import synchronized
|
|
|
|
|
|
def zip_(*args: Observable[Any]) -> Observable[Tuple[Any, ...]]:
|
|
"""Merges the specified observable sequences into one observable
|
|
sequence by creating a tuple whenever all of the
|
|
observable sequences have produced an element at a corresponding
|
|
index.
|
|
|
|
Example:
|
|
>>> res = zip(obs1, obs2)
|
|
|
|
Args:
|
|
args: Observable sources to zip.
|
|
|
|
Returns:
|
|
An observable sequence containing the result of combining
|
|
elements of the sources as tuple.
|
|
"""
|
|
|
|
sources = list(args)
|
|
|
|
def subscribe(
|
|
observer: abc.ObserverBase[Any], scheduler: Optional[abc.SchedulerBase] = None
|
|
) -> CompositeDisposable:
|
|
n = len(sources)
|
|
queues: List[List[Any]] = [[] for _ in range(n)]
|
|
lock = RLock()
|
|
is_completed = [False] * n
|
|
|
|
@synchronized(lock)
|
|
def next_(i: int) -> None:
|
|
if all(len(q) for q in queues):
|
|
try:
|
|
queued_values = [x.pop(0) for x in queues]
|
|
res = tuple(queued_values)
|
|
except Exception as ex: # pylint: disable=broad-except
|
|
observer.on_error(ex)
|
|
return
|
|
|
|
observer.on_next(res)
|
|
|
|
# after sending the zipped values, complete the observer if at least one
|
|
# upstream observable is completed and its queue has length zero
|
|
if any(
|
|
(
|
|
done
|
|
for queue, done in zip(queues, is_completed)
|
|
if len(queue) == 0
|
|
)
|
|
):
|
|
observer.on_completed()
|
|
|
|
def completed(i: int) -> None:
|
|
is_completed[i] = True
|
|
if len(queues[i]) == 0:
|
|
observer.on_completed()
|
|
|
|
subscriptions: List[Optional[abc.DisposableBase]] = [None] * n
|
|
|
|
def func(i: int) -> None:
|
|
source: Observable[Any] = sources[i]
|
|
if isinstance(source, Future):
|
|
source = from_future(source)
|
|
|
|
sad = SingleAssignmentDisposable()
|
|
|
|
def on_next(x: Any) -> None:
|
|
queues[i].append(x)
|
|
next_(i)
|
|
|
|
sad.disposable = source.subscribe(
|
|
on_next, observer.on_error, lambda: completed(i), scheduler=scheduler
|
|
)
|
|
subscriptions[i] = sad
|
|
|
|
for idx in range(n):
|
|
func(idx)
|
|
return CompositeDisposable(subscriptions)
|
|
|
|
return Observable(subscribe)
|
|
|
|
|
|
__all__ = ["zip_"]
|