Add switch_map operator and equivalent starred and indexed

This draws from the definition in rxjs and maintains parity with the
map operator and its variants

Author: giff-h

reactivex/operators/__init__.py

@@ -26,6 +26,7 @@
     Observable,
     abc,
     compose,
+    of,
     typing,
 )
 from reactivex.internal.basic import identity
@@ -3325,6 +3326,202 @@ def switch_latest() -> Callable[
     return switch_latest_()
 
 
+def switch_map(
+    mapper: Optional[Mapper[_T1, Observable[_T2]]] = None
+) -> Callable[[Observable[_T1]], Observable[_T2]]:
+    """
+    The switch_map operator.
+
+    Project each element of an observable sequence into a new observable.
+
+    .. marble::
+        :alt: switch_map
+
+        ---1---2---3--->
+        [   switch_map(i: of(i, i ** 2, i ** 3))    ]
+        ---1---1---1---2---4---8---3---9---27--->
+
+    Example:
+        >>> switch_map(lambda value: of(value, value // 2))
+
+    Args:
+        mapper: A transform function to apply to each source element.
+
+    Returns:
+        A partially applied operator function that takes an observable
+        source and returns an observable sequence whose elements are
+        each element of the result of invoking the transform function
+        on each element of the source.
+    """
+    from ._switch_map import switch_map_
+
+    return switch_map_(mapper)
+
+
+def switch_map_indexed(
+    mapper_indexed: Optional[MapperIndexed[_T1, Observable[_T2]]] = None
+) -> Callable[[Observable[_T1]], Observable[_T2]]:
+    """
+    The switch_map_indexed operator.
+
+    Project each element of an observable sequence into a new observable
+    by incorporating the element's index.
+
+    .. marble::
+        :alt: switch_map_indexed
+
+        ---1-----------2-----------3----------->
+        [   switch_map_indexed(i,id: of(i, i ** 2, i + id))    ]
+        ---1---1---1---2---4---3---3---9---5--->
+
+    Example:
+        >>> switch_map_indexed(lambda value, index: of(value, value // 2))
+
+    Args:
+        mapper_indexed: A transform function to apply to each source
+            element. The second parameter of the function represents
+            the index of the source element.
+
+    Returns:
+        A partially applied operator function that takes an observable
+        source and returns an observable sequence whose elements are
+        each element of the result of invoking the transform function
+        on each element of the source.
+    """
+    from ._switch_map import switch_map_indexed_
+
+    return switch_map_indexed_(mapper_indexed)
+
+
+@overload
+def switch_starmap(
+    mapper: Callable[[_A, _B], Observable[_T]]
+) -> Callable[[Observable[Tuple[_A, _B]]], Observable[_T]]:
+    ...
+
+
+@overload
+def switch_starmap(
+    mapper: Callable[[_A, _B, _C], Observable[_T]]
+) -> Callable[[Observable[Tuple[_A, _B, _C]]], Observable[_T]]:
+    ...
+
+
+@overload
+def switch_starmap(
+    mapper: Callable[[_A, _B, _C, _D], Observable[_T]]
+) -> Callable[[Observable[Tuple[_A, _B, _C, _D]]], Observable[_T]]:
+    ...
+
+
+def switch_starmap(
+    mapper: Optional[Callable[..., Observable[Any]]] = None
+) -> Callable[[Observable[Any]], Observable[Any]]:
+    """The switch_starmap operator.
+
+    Unpack arguments grouped as tuple elements of an observable sequence
+    and return an observable sequence whose values are each element of
+    the observable returned by invoking the mapper function with star
+    applied on unpacked elements as positional arguments.
+
+    Use instead of `switch_map()` when the the arguments to the mapper is
+    grouped as tuples and the mapper function takes multiple arguments.
+
+    .. marble::
+        :alt: switch_starmap
+
+        ----1,2-------3,4---------|
+        [   switch_starmap(of)    ]
+        ----1----2----3----4------|
+
+    Example:
+        >>> switch_starmap(lambda x, y: of(x + y, x * y))
+
+    Args:
+        mapper: A transform function to invoke with unpacked elements
+            as arguments.
+
+    Returns:
+        An operator function that takes an observable source and returns
+        an observable sequence whose values are each element of the
+        observable returned by invoking the mapper function with the
+        unpacked elements of the source.
+    """
+
+    if mapper is None:
+        mapper = of
+
+    def starred(values: Tuple[Any, ...]) -> Observable[Any]:
+        return mapper(*values)
+
+    return compose(switch_map(starred))
+
+
+@overload
+def switch_starmap_indexed(
+    mapper: Callable[[_A, int], Observable[_T]]
+) -> Callable[[Observable[_A]], Observable[_T]]:
+    ...
+
+
+@overload
+def switch_starmap_indexed(
+    mapper: Callable[[_A, _B, int], Observable[_T]]
+) -> Callable[[Observable[Tuple[_A, _B]]], Observable[_T]]:
+    ...
+
+
+@overload
+def switch_starmap_indexed(
+    mapper: Callable[[_A, _B, _C, int], Observable[_T]]
+) -> Callable[[Observable[Tuple[_A, _B, _C]]], Observable[_T]]:
+    ...
+
+
+@overload
+def switch_starmap_indexed(
+    mapper: Callable[[_A, _B, _C, _D, int], Observable[_T]]
+) -> Callable[[Observable[Tuple[_A, _B, _C, _D]]], Observable[_T]]:
+    ...
+
+
+def switch_starmap_indexed(
+    mapper: Optional[Callable[..., Observable[Any]]] = None
+) -> Callable[[Observable[Any]], Observable[Any]]:
+    """Variant of :func:`switch_starmap` which accepts an indexed mapper.
+
+    .. marble::
+        :alt: switch_starmap_indexed
+
+        ------1,2----------3,4-----------|
+        [   switch_starmap_indexed(of)   ]
+        ------1---2---0----3---4---1-----|
+
+    Example:
+        >>> switch_starmap_indexed(lambda x, y, i: of(x + y + i, x * y - i))
+
+    Args:
+        mapper: A transform function to invoke with unpacked elements
+            as arguments.
+
+    Returns:
+        An operator function that takes an observable source and returns
+        an observable sequence whose values are each element of the
+        observable returned by invoking the mapper function with the
+        unpacked elements of the source.
+    """
+    from ._switch_map import switch_map_
+
+    if mapper is None:
+        return compose(of)
+
+    def starred(values: Tuple[Any, ...]) -> Observable[Any]:
+        assert mapper  # mypy is paranoid
+        return mapper(*values)
+
+    return compose(switch_map_(starred))
+
+
 def take(count: int) -> Callable[[Observable[_T]], Observable[_T]]:
     """Returns a specified number of contiguous elements from the start
     of an observable sequence.
@@ -4272,6 +4469,10 @@ def zip_with_iterable(
     "subscribe_on",
     "sum",
     "switch_latest",
+    "switch_map",
+    "switch_map_indexed",
+    "switch_starmap",
+    "switch_starmap_indexed",
     "take",
     "take_last",
     "take_last_buffer",

reactivex/operators/_switch_map.py

@@ -0,0 +1,96 @@
+from typing import Callable, Optional, TypeVar, cast
+
+from reactivex import Observable, abc, compose, of
+from reactivex import operators as ops
+from reactivex.internal import infinite
+from reactivex.typing import Mapper, MapperIndexed
+
+_T1 = TypeVar("_T1")
+_T2 = TypeVar("_T2")
+
+
+def switch_map_(
+    mapper: Optional[Mapper[_T1, Observable[_T2]]] = None
+) -> Callable[[Observable[_T1]], Observable[_T2]]:
+    _mapper = mapper or cast(Mapper[_T1, Observable[_T2]], of)
+
+    def switch_map(source: Observable[_T1]) -> Observable[_T2]:
+        """
+        Partially applied switch_map operator.
+
+        Project each element of an observable sequence into a new observable
+        by incorporating the element's index.
+
+        Example:
+            >>> switch_map(source)
+
+        Args:
+            source: The observable source to transform.
+
+        Returns:
+            Returns an observable sequence whose elements are each element of
+            the result of invoking the transform function on each element of the
+            source.
+        """
+
+        inner_observer: Optional[abc.DisposableBase] = None
+        is_complete = False
+
+        def subscribe(
+            obv: abc.ObserverBase[_T2], scheduler: Optional[abc.SchedulerBase] = None
+        ) -> abc.DisposableBase:
+            nonlocal inner_observer, is_complete
+
+            def check_complete() -> None:
+                if is_complete and not inner_observer:
+                    obv.on_completed()
+
+            def on_next(value: _T1) -> None:
+                nonlocal inner_observer
+
+                def inner_complete():
+                    nonlocal inner_observer
+                    inner_observer = None
+                    check_complete()
+
+                if inner_observer:
+                    inner_observer.dispose()
+
+                try:
+                    mapped_source = _mapper(value)
+                except Exception as err:
+                    obv.on_error(err)
+                else:
+                    inner_observer = mapped_source.subscribe(
+                        obv.on_next, obv.on_error, inner_complete, scheduler=scheduler
+                    )
+
+            def source_complete():
+                nonlocal is_complete
+                is_complete = True
+                check_complete()
+
+            return source.subscribe(
+                on_next, obv.on_error, source_complete, scheduler=scheduler
+            )
+
+        return Observable(subscribe)
+
+    return switch_map
+
+
+def switch_map_indexed_(
+    mapper_indexed: Optional[MapperIndexed[_T1, Observable[_T2]]] = None
+) -> Callable[[Observable[_T1]], Observable[_T2]]:
+    def _of(value: _T1, _: int) -> Observable[_T2]:
+        return of(cast(_T2, value))
+
+    _mapper_indexed = mapper_indexed or cast(MapperIndexed[_T1, Observable[_T2]], _of)
+
+    return compose(
+        ops.zip_with_iterable(infinite()),
+        ops.switch_starmap_indexed(_mapper_indexed),
+    )
+
+
+__all__ = ["switch_map_", "switch_map_indexed_"]