-
Notifications
You must be signed in to change notification settings - Fork 623
/
filter.kt
42 lines (37 loc) · 1.56 KB
/
filter.kt
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
package io.kotest.property.arbitrary
import io.kotest.property.Arb
import io.kotest.property.PropertyTesting
import io.kotest.property.RTree
import io.kotest.property.RandomSource
import io.kotest.property.Sample
import io.kotest.property.filter
/**
* Returns a new [Arb] which takes its elements from the receiver and filters them using the supplied
* predicate. This gen will continue to request elements from the underlying gen until one satisfies
* the predicate.
*/
fun <A> Arb<A>.filter(predicate: (A) -> Boolean): Arb<A> = trampoline { sampleA ->
object : Arb<A>() {
override fun edgecase(rs: RandomSource): A? =
sequenceOf(sampleA.value)
.plus(generateSequence { this@filter.edgecase(rs) })
.take(PropertyTesting.maxFilterAttempts)
.filter(predicate)
.firstOrNull()
override fun sample(rs: RandomSource): Sample<A> {
val sample = sequenceOf(sampleA).plus(this@filter.samples(rs)).filter { predicate(it.value) }.first()
return Sample(sample.value, sample.shrinks.filter(predicate) ?: RTree({ sample.value }))
}
}
}
/**
* @return a new [Arb] by filtering this arbs output by the negated function [f]
*/
fun <A> Arb<A>.filterNot(f: (A) -> Boolean): Arb<A> = filter { !f(it) }
/**
* Create a new [Arb] by keeping only instances of B generated by this gen.
* This is useful if you have a type hierarchy and only want to retain
* a particular subtype.
*/
@Suppress("UNCHECKED_CAST")
inline fun <A, reified B : A> Arb<A>.filterIsInstance(): Arb<B> = filter { it is B }.map { it as B }