Add StateFlow<T>.onSubscription

[xit0c]

Use case

The current SharedFlow<T>.onSubscription implementation returns a SharedFlow<T>, which means losing type information when calling it from a StateFlow<T>. It would be nice to have a StateFlow<T>.onSubscription counterpart for StateFlow<T>, thus avoiding the aforementioned issue.

The Shape of the API

Unless there are implementation details I'm unaware of, the API could be copy-pasted from the SharedFlow<T>.onSubscription implementation and adapted to work with StateFlow<T>:

public fun <T> StateFlow<T>.onSubscription(action: suspend FlowCollector<T>.() -> Unit): StateFlow<T> =
    SubscribedStateFlow(this, action)

@OptIn(ExperimentalForInheritanceCoroutinesApi::class)
private class SubscribedStateFlow<T>(
    private val stateFlow: StateFlow<T>,
    private val action: suspend FlowCollector<T>.() -> Unit
) : StateFlow<T> by stateFlow {
    override suspend fun collect(collector: FlowCollector<T>) =
        stateFlow.collect(SubscribedFlowCollector(collector, action))
}

Prior Art

• SharedFlow.onSubscription
• SubscribedSharedFlow