diff --git a/src/rapidfuzz/process_py.py b/src/rapidfuzz/process_py.py index dcee33aa..d3082c7b 100644 --- a/src/rapidfuzz/process_py.py +++ b/src/rapidfuzz/process_py.py @@ -3,6 +3,9 @@ from __future__ import annotations import heapq +import operator +import struct +import sys from rapidfuzz._utils import ScorerFlag, is_none, setupPandas from rapidfuzz.fuzz import WRatio, ratio @@ -14,8 +17,46 @@ def _get_scorer_flags_py(scorer, scorer_kwargs): params = getattr(scorer, "_RF_ScorerPy", None) if params is not None: flags = params["get_scorer_flags"](**scorer_kwargs) - return (flags["worst_score"], flags["optimal_score"]) - return (0, 100) + return (flags["worst_score"], flags["optimal_score"], flags["flags"]) + return (0, 100, None) + + +def _validate_score_cutoff(score_cutoff, worst_score, optimal_score, scorer_flags): + if score_cutoff is None or scorer_flags is None: + return score_cutoff + + if scorer_flags & ScorerFlag.RESULT_F64: + if not hasattr(score_cutoff, "__float__"): + error_message = f"must be real number, not {type(score_cutoff).__name__}" + raise TypeError(error_message) + score_cutoff = struct.unpack("f", struct.pack("f", float(score_cutoff)))[0] + worst_score = struct.unpack("f", struct.pack("f", float(worst_score)))[0] + optimal_score = struct.unpack("f", struct.pack("f", float(optimal_score)))[0] + else: + try: + score_cutoff = operator.index(score_cutoff) + except TypeError: + error_message = "an integer is required" + raise TypeError(error_message) from None + if scorer_flags & ScorerFlag.RESULT_SIZE_T: + if score_cutoff < 0: + error_message = "can't convert negative value to uint64_t" + raise OverflowError(error_message) + size_t_max = 2 * sys.maxsize + 1 + if score_cutoff > size_t_max: + error_message = "Python int too large to convert to C unsigned long" + raise OverflowError(error_message) + if optimal_score > worst_score: + optimal_score = size_t_max + else: + worst_score = size_t_max + + lowest_score = min(worst_score, optimal_score) + highest_score = max(worst_score, optimal_score) + if score_cutoff < lowest_score or score_cutoff > highest_score: + error_message = f"score_cutoff has to be in the range of {lowest_score} - {highest_score}" + raise TypeError(error_message) + return score_cutoff def extract_iter( @@ -85,9 +126,8 @@ def extract_iter( * The `key of choice` when choices is a mapping like a dict, or a pandas Series """ - _ = score_hint scorer_kwargs = scorer_kwargs or {} - worst_score, optimal_score = _get_scorer_flags_py(scorer, scorer_kwargs) + worst_score, optimal_score, scorer_flags = _get_scorer_flags_py(scorer, scorer_kwargs) lowest_score_worst = optimal_score > worst_score setupPandas() @@ -95,6 +135,9 @@ def extract_iter( if is_none(query): return + score_cutoff = _validate_score_cutoff(score_cutoff, worst_score, optimal_score, scorer_flags) + _validate_score_cutoff(score_hint, worst_score, optimal_score, scorer_flags) + if score_cutoff is None: score_cutoff = worst_score @@ -251,9 +294,8 @@ def extractOne( None """ - _ = score_hint scorer_kwargs = scorer_kwargs or {} - worst_score, optimal_score = _get_scorer_flags_py(scorer, scorer_kwargs) + worst_score, optimal_score, scorer_flags = _get_scorer_flags_py(scorer, scorer_kwargs) lowest_score_worst = optimal_score > worst_score setupPandas() @@ -261,6 +303,9 @@ def extractOne( if is_none(query): return None + score_cutoff = _validate_score_cutoff(score_cutoff, worst_score, optimal_score, scorer_flags) + _validate_score_cutoff(score_hint, worst_score, optimal_score, scorer_flags) + if score_cutoff is None: score_cutoff = worst_score @@ -376,7 +421,7 @@ def extract( """ scorer_kwargs = scorer_kwargs or {} - worst_score, optimal_score = _get_scorer_flags_py(scorer, scorer_kwargs) + worst_score, optimal_score, _ = _get_scorer_flags_py(scorer, scorer_kwargs) lowest_score_worst = optimal_score > worst_score if limit == 1: @@ -514,8 +559,12 @@ def cdist( """ import numpy as np - _ = workers, score_hint + _ = workers scorer_kwargs = scorer_kwargs or {} + if getattr(scorer, "_RF_OriginalScorer", None) is scorer: + worst_score, optimal_score, scorer_flags = _get_scorer_flags_py(scorer, scorer_kwargs) + score_cutoff = _validate_score_cutoff(score_cutoff, worst_score, optimal_score, scorer_flags) + _validate_score_cutoff(score_hint, worst_score, optimal_score, scorer_flags) dtype = _dtype_to_type_num(dtype, scorer, scorer_kwargs) results = np.zeros((len(queries), len(choices)), dtype=dtype) @@ -649,8 +698,12 @@ def cpdist( error_message = "Length of queries and choices must be the same!" raise ValueError(error_message) - _ = workers, score_hint + _ = workers scorer_kwargs = scorer_kwargs or {} + if getattr(scorer, "_RF_OriginalScorer", None) is scorer: + worst_score, optimal_score, scorer_flags = _get_scorer_flags_py(scorer, scorer_kwargs) + score_cutoff = _validate_score_cutoff(score_cutoff, worst_score, optimal_score, scorer_flags) + _validate_score_cutoff(score_hint, worst_score, optimal_score, scorer_flags) dtype = _dtype_to_type_num(dtype, scorer, scorer_kwargs) results = np.zeros((len_queries,), dtype=dtype) diff --git a/tests/test_process.py b/tests/test_process.py index e0ecd7bd..e71c3784 100644 --- a/tests/test_process.py +++ b/tests/test_process.py @@ -97,6 +97,54 @@ def test_extractOne_exceptions(): process_py.extractOne("", {1: 1}) +def call_process(process_impl, function, scorer, **kwargs): + if function == "extract_iter": + return list(process_impl.extract_iter("a", ["a"], scorer=scorer, **kwargs)) + if function in {"cdist", "cpdist"}: + return getattr(process_impl, function)(["a"], ["a"], scorer=scorer, **kwargs) + return getattr(process_impl, function)("a", ["a"], scorer=scorer, **kwargs) + + +@pytest.mark.parametrize("process_impl", [process_cpp, process_py]) +@pytest.mark.parametrize("function", ["extractOne", "extract", "extract_iter", "cdist", "cpdist"]) +@pytest.mark.parametrize("argument", ["score_cutoff", "score_hint"]) +@pytest.mark.parametrize( + ("scorer", "value", "score_range"), + [ + (fuzz.ratio, -0.1, "0.0 - 100.0"), + (fuzz.ratio, 100.1, "0.0 - 100.0"), + (Levenshtein.normalized_similarity, -0.1, "0.0 - 1.0"), + (Levenshtein.normalized_similarity, 1.1, "0.0 - 1.0"), + ], +) +def test_process_score_range(process_impl, function, argument, scorer, value, score_range): + with pytest.raises(TypeError, match=f"score_cutoff has to be in the range of {score_range}"): + call_process(process_impl, function, scorer, **{argument: value}) + + +@pytest.mark.parametrize("process_impl", [process_cpp, process_py]) +def test_process_score_cutoff_float32_rounding(process_impl): + result = process_impl.extractOne("a", ["a"], scorer=fuzz.ratio, score_cutoff=100.000001) + assert result == ("a", 100.0, 0) + + +@pytest.mark.parametrize("process_impl", [process_cpp, process_py]) +@pytest.mark.parametrize( + ("scorer", "value", "error_message"), + [(fuzz.ratio, "50", "must be real number, not str"), (Levenshtein.distance, "1", "an integer is required")], +) +def test_process_score_cutoff_rejects_string(process_impl, scorer, value, error_message): + with pytest.raises(TypeError, match=error_message): + process_impl.extractOne("a", ["a"], scorer=scorer, score_cutoff=value) + + +@pytest.mark.parametrize("process_impl", [process_cpp, process_py]) +@pytest.mark.parametrize("function", ["extractOne", "extract", "extract_iter", "cdist", "cpdist"]) +def test_process_distance_score_hint_range(process_impl, function): + with pytest.raises(OverflowError, match="can't convert negative value to uint64_t"): + call_process(process_impl, function, Levenshtein.distance, score_hint=-1) + + def test_extract_exceptions(): with pytest.raises(TypeError): process_cpp.extract()