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}