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

[dkhalanskyjb]

Sure, this makes sense, thanks! Would you like to open a PR with this code (+ documentation + tests)?

[xit0c]

I'd be glad to! I hope to deliver the changes in the next few days.

[xit0c]

Done!

The documentation is copy-pasted from SharedFlow.onSubscription, I've only replaced the word "shared" with "state" in the first line.

About the tests, I saw that there's already testOnSubscriptionWithException (which now is using the new onSubscription), is it enough?

[dkhalanskyjb]

Good job, thanks! Let's continue under #4380, as this starts to be about the implementation and not the issue itself.