commonMain.net.ormr.fuzzywuzzy.Extractor.kt Maven / Gradle / Ivy

Go to download
package net.ormr.fuzzywuzzy

import net.ormr.fuzzywuzzy.algorithms.buildMutableList
import net.ormr.fuzzywuzzy.algorithms.findTopKHeap
import net.ormr.fuzzywuzzy.model.BoundExtractedResult
import net.ormr.fuzzywuzzy.model.ExtractedResult

internal class Extractor private constructor(val cutoff: Int) {
    /**
     * Returns the list of choices with their associated scores of similarity in a list
     * of [ExtractedResult].
     */
    fun extractWithoutOrder(
        query: String,
        choices: Collection,
        scoreFunc: Applicable,
    ): MutableList = buildMutableList {
        for ((index, s) in choices.withIndex()) {
            val score = scoreFunc(query, s)

            if (score >= cutoff) {
                add(ExtractedResult(s, score, index))
            }
        }
    }

    /**
     * Returns the list of choices with their associated scores of similarity in a list
     * of [ExtractedResult].
     */
    fun  extractWithoutOrder(
        query: String,
        choices: Collection,
        toString: ToStringFunction,
        scoreFunc: Applicable,
    ): MutableList> = buildMutableList {
        for ((index, t) in choices.withIndex()) {
            val s = toString(t)
            val score = scoreFunc(query, s)

            if (score >= cutoff) {
                add(BoundExtractedResult(t, s, score, index))
            }
        }
    }

    /**
     * Find the single best match above a score in a list of choices.
     */
    fun extractOne(
        query: String,
        choices: Collection,
        scoreFunc: Applicable,
    ): ExtractedResult = extractWithoutOrder(query, choices, scoreFunc).max()

    /**
     * Find the single best match above a score in a list of choices.
     */
    fun  extractOne(
        query: String,
        choices: Collection,
        toString: ToStringFunction,
        scoreFunc: Applicable,
    ): BoundExtractedResult = extractWithoutOrder(query, choices, toString, scoreFunc).max()

    /**
     * Creates a **sorted** list of [ExtractedResult]  which contain the
     * top @param limit most similar choices.
     */
    fun extractTop(query: String, choices: Collection, func: Applicable): List =
        extractWithoutOrder(query, choices, func).apply { sortWith(reverseOrder()) }

    /**
     * Creates a **sorted** list of [ExtractedResult]  which contain the
     * top @param limit most similar choices.
     */
    fun  extractTop(
        query: String,
        choices: Collection,
        toString: ToStringFunction,
        scoreFunc: Applicable,
    ): List> = extractWithoutOrder(query, choices, toString, scoreFunc)
        .apply { sortWith(reverseOrder()) }

    /**
     * Creates a **sorted** list of [ExtractedResult] which contain the
     * top @param limit most similar choices.
     */
    fun extractTop(
        query: String,
        choices: Collection,
        scoreFunc: Applicable,
        limit: Int,
    ): List = extractWithoutOrder(query, choices, scoreFunc)
        .findTopKHeap(limit)
        .asReversed()

    /**
     * Creates a **sorted** list of [ExtractedResult] which contain the
     * top @param limit most similar choices.
     */
    fun  extractTop(
        query: String,
        choices: Collection,
        toString: ToStringFunction,
        scoreFunc: Applicable,
        limit: Int,
    ): List> = extractWithoutOrder(query, choices, toString, scoreFunc)
        .findTopKHeap(limit)
        .asReversed()

    // sequence

    fun extractWithoutOrder(
        query: String,
        choices: Sequence,
        scoreFunc: Applicable,
    ): Sequence = choices.mapIndexed { index, item ->
        val score = scoreFunc(query, item)

        ExtractedResult(item, score, index)
    }.filter { it.score >= cutoff }

    fun  extractWithoutOrder(
        query: String,
        choices: Sequence,
        toString: ToStringFunction,
        scoreFunc: Applicable,
    ): Sequence> = choices.mapIndexed { index, item ->
        val string = toString(item)
        val score = scoreFunc(query, string)

        BoundExtractedResult(item, string, score, index)
    }.filter { it.score >= cutoff }

    fun extractOne(
        query: String,
        choices: Sequence,
        scoreFunc: Applicable,
    ): ExtractedResult = extractWithoutOrder(query, choices, scoreFunc).max()

    fun  extractOne(
        query: String,
        choices: Sequence,
        toString: ToStringFunction,
        scoreFunc: Applicable,
    ): BoundExtractedResult = extractWithoutOrder(query, choices, toString, scoreFunc).max()

    fun extractTop(query: String, choices: Sequence, func: Applicable): Sequence =
        extractWithoutOrder(query, choices, func).sortedWith(reverseOrder())

    fun  extractTop(
        query: String,
        choices: Sequence,
        toString: ToStringFunction,
        scoreFunc: Applicable,
    ): Sequence> = extractWithoutOrder(query, choices, toString, scoreFunc)
        .sortedWith(reverseOrder())

    fun extractTop(
        query: String,
        choices: Sequence,
        scoreFunc: Applicable,
        limit: Int,
    ): List = extractWithoutOrder(query, choices, scoreFunc)
        .findTopKHeap(limit)
        .toList()
        .asReversed()

    fun  extractTop(
        query: String,
        choices: Sequence,
        toString: ToStringFunction,
        scoreFunc: Applicable,
        limit: Int,
    ): List> = extractWithoutOrder(query, choices, toString, scoreFunc)
        .findTopKHeap(limit)
        .toList()
        .asReversed()

    companion object {
        val NO_CUTOFF: Extractor = Extractor(cutoff = 0)

        fun of(cutoff: Int): Extractor = when (cutoff) {
            0 -> NO_CUTOFF
            else -> Extractor(cutoff)
        }
    }
}