Open
Description
takeXxx
share a common logic, have you considered collectWhile
?
Here a POC
@FlowPreview
public fun <T> Flow<T>.take(count: Int): Flow<T> {
if (count == 0) return emptyFlow()
require(count > 0) { "Take count should be non negative, but had $count" }
return flow {
var consumed = 0
collectWhile { value ->
if (consumed != count) {
emit(value)
consumed++
true
} else false
}
}
}
public suspend fun <T> Flow<T>.collectWhile(action: suspend (value: T) -> Boolean) {
var cancellationException: CancellationException? = null
try {
collect { value ->
val loop = action(value)
if (!loop) {
val exception = CancellationException()
cancellationException = exception
throw exception
}
}
} catch (e: CancellationException) {
if (e !== cancellationException) throw e
}
}