from typing import Any, Callable, List, Optional, TypeVar from reactivex import Observable, abc, typing from reactivex.disposable import ( CompositeDisposable, SerialDisposable, SingleAssignmentDisposable, ) from reactivex.scheduler import ImmediateScheduler _T = TypeVar("_T") def expand_( mapper: typing.Mapper[_T, Observable[_T]] ) -> Callable[[Observable[_T]], Observable[_T]]: def expand(source: Observable[_T]) -> Observable[_T]: """Expands an observable sequence by recursively invoking mapper. Args: source: Source obserable to expand. Returns: An observable sequence containing all the elements produced by the recursive expansion. """ def subscribe( observer: abc.ObserverBase[_T], scheduler: Optional[abc.SchedulerBase] = None, ) -> abc.DisposableBase: scheduler = scheduler or ImmediateScheduler.singleton() queue: List[Observable[_T]] = [] m = SerialDisposable() d = CompositeDisposable(m) active_count = 0 is_acquired = False def ensure_active(): nonlocal is_acquired is_owner = False if queue: is_owner = not is_acquired is_acquired = True def action(scheduler: abc.SchedulerBase, state: Any = None): nonlocal is_acquired, active_count if queue: work = queue.pop(0) else: is_acquired = False return sad = SingleAssignmentDisposable() d.add(sad) def on_next(value: _T) -> None: nonlocal active_count observer.on_next(value) result = None try: result = mapper(value) except Exception as ex: observer.on_error(ex) return queue.append(result) active_count += 1 ensure_active() def on_complete() -> None: nonlocal active_count d.remove(sad) active_count -= 1 if active_count == 0: observer.on_completed() sad.disposable = work.subscribe( on_next, observer.on_error, on_complete, scheduler=scheduler ) m.disposable = scheduler.schedule(action) if is_owner: m.disposable = scheduler.schedule(action) queue.append(source) active_count += 1 ensure_active() return d return Observable(subscribe) return expand __all__ = ["expand_"]