from typing import List, Optional, TypeVar from reactivex import abc from reactivex.disposable import CompositeDisposable, Disposable from .observable import Observable _T = TypeVar("_T") class ConnectableObservable(Observable[_T]): """Represents an observable that can be connected and disconnected.""" def __init__(self, source: abc.ObservableBase[_T], subject: abc.SubjectBase[_T]): self.subject = subject self.has_subscription = False self.subscription: Optional[abc.DisposableBase] = None self.source = source super().__init__() def _subscribe_core( self, observer: abc.ObserverBase[_T], scheduler: Optional[abc.SchedulerBase] = None, ) -> abc.DisposableBase: return self.subject.subscribe(observer, scheduler=scheduler) def connect( self, scheduler: Optional[abc.SchedulerBase] = None ) -> Optional[abc.DisposableBase]: """Connects the observable.""" if not self.has_subscription: self.has_subscription = True def dispose() -> None: self.has_subscription = False subscription = self.source.subscribe(self.subject, scheduler=scheduler) self.subscription = CompositeDisposable(subscription, Disposable(dispose)) return self.subscription def auto_connect(self, subscriber_count: int = 1) -> Observable[_T]: """Returns an observable sequence that stays connected to the source indefinitely to the observable sequence. Providing a subscriber_count will cause it to connect() after that many subscriptions occur. A subscriber_count of 0 will result in emissions firing immediately without waiting for subscribers. """ connectable_subscription: List[Optional[abc.DisposableBase]] = [None] count = [0] source = self is_connected = [False] if subscriber_count == 0: connectable_subscription[0] = source.connect() is_connected[0] = True def subscribe( observer: abc.ObserverBase[_T], scheduler: Optional[abc.SchedulerBase] = None, ) -> abc.DisposableBase: count[0] += 1 should_connect = count[0] == subscriber_count and not is_connected[0] subscription = source.subscribe(observer) if should_connect: connectable_subscription[0] = source.connect(scheduler) is_connected[0] = True def dispose() -> None: subscription.dispose() count[0] -= 1 is_connected[0] = False return Disposable(dispose) return Observable(subscribe)