Skip to main content

matchr/
python.rs

1use crate::{
2    batch_best_match as bbm_fn, best_match as bm_fn, combined_score as cs_fn,
3    jaro_winkler as jw_fn, levenshtein as lev_fn, rank_matches as rm_fn,
4    token_set_ratio as tset_fn, token_sort_ratio as tsort_fn, trigram_similarity as tri_fn,
5};
6use pyo3::prelude::*;
7
8#[pyfunction]
9fn levenshtein(a: &str, b: &str) -> usize {
10    lev_fn(a, b)
11}
12
13#[pyfunction]
14fn jaro_winkler(a: &str, b: &str) -> f64 {
15    jw_fn(a, b)
16}
17
18#[pyfunction]
19fn trigram_similarity(a: &str, b: &str) -> f64 {
20    tri_fn(a, b)
21}
22
23#[pyfunction]
24fn combined_score(a: &str, b: &str) -> f64 {
25    cs_fn(a, b)
26}
27
28#[pyfunction]
29fn token_sort_ratio(a: &str, b: &str) -> f64 {
30    tsort_fn(a, b)
31}
32
33#[pyfunction]
34fn token_set_ratio(a: &str, b: &str) -> f64 {
35    tset_fn(a, b)
36}
37
38#[pyfunction]
39#[pyo3(signature = (query, candidates, threshold=None))]
40fn best_match(
41    query: &str,
42    candidates: Vec<String>,
43    threshold: Option<f64>,
44) -> Option<(String, f64)> {
45    let refs: Vec<&str> = candidates.iter().map(|s| s.as_str()).collect();
46    let min = threshold.unwrap_or(0.0);
47    bm_fn(query, &refs)
48        .filter(|(_, score)| *score >= min)
49        .map(|(s, score)| (s.to_string(), score))
50}
51
52#[pyfunction]
53#[pyo3(signature = (query, candidates, threshold=None))]
54fn rank_matches(
55    query: &str,
56    candidates: Vec<String>,
57    threshold: Option<f64>,
58) -> Vec<(String, f64)> {
59    let refs: Vec<&str> = candidates.iter().map(|s| s.as_str()).collect();
60    let min = threshold.unwrap_or(0.0);
61    rm_fn(query, &refs)
62        .into_iter()
63        .filter(|(_, score)| *score >= min)
64        .map(|(s, score)| (s.to_string(), score))
65        .collect()
66}
67
68#[pyfunction]
69#[pyo3(signature = (queries, candidates, threshold=None))]
70fn batch_best_match(
71    queries: Vec<String>,
72    candidates: Vec<String>,
73    threshold: Option<f64>,
74) -> Vec<Option<(String, f64)>> {
75    let q_refs: Vec<&str> = queries.iter().map(|s| s.as_str()).collect();
76    let c_refs: Vec<&str> = candidates.iter().map(|s| s.as_str()).collect();
77    let min = threshold.unwrap_or(0.0);
78
79    bbm_fn(&q_refs, &c_refs)
80        .into_iter()
81        .map(|opt| {
82            opt.filter(|(_, score)| *score >= min)
83                .map(|(s, score)| (s.to_string(), score))
84        })
85        .collect()
86}
87
88#[pymodule]
89fn matchr(m: &Bound<'_, PyModule>) -> PyResult<()> {
90    m.add_function(wrap_pyfunction!(levenshtein, m)?)?;
91    m.add_function(wrap_pyfunction!(jaro_winkler, m)?)?;
92    m.add_function(wrap_pyfunction!(trigram_similarity, m)?)?;
93    m.add_function(wrap_pyfunction!(combined_score, m)?)?;
94    m.add_function(wrap_pyfunction!(token_sort_ratio, m)?)?;
95    m.add_function(wrap_pyfunction!(token_set_ratio, m)?)?;
96    m.add_function(wrap_pyfunction!(best_match, m)?)?;
97    m.add_function(wrap_pyfunction!(rank_matches, m)?)?;
98    m.add_function(wrap_pyfunction!(batch_best_match, m)?)?;
99    Ok(())
100}