# SPDX-License-Identifier: MIT
# Copyright (C) 2022 Max Bachmann
from __future__ import annotations
from rapidfuzz.fuzz import ratio
from rapidfuzz.process_cpp_impl import (
FLOAT32 as _FLOAT32,
FLOAT64 as _FLOAT64,
INT8 as _INT8,
INT16 as _INT16,
INT32 as _INT32,
INT64 as _INT64,
UINT8 as _UINT8,
UINT16 as _UINT16,
UINT32 as _UINT32,
UINT64 as _UINT64,
cdist as _cdist,
cpdist as _cpdist,
extract,
extract_iter,
extractOne,
)
__all__ = ["extract", "extract_iter", "extractOne", "cdist", "cpdist"]
def _dtype_to_type_num(dtype):
import numpy as np
if dtype is None:
return None
dtype = np.dtype(dtype)
if dtype == np.int32:
return _INT32
if dtype == np.int8:
return _INT8
if dtype == np.int16:
return _INT16
if dtype == np.int64:
return _INT64
if dtype == np.uint8:
return _UINT8
if dtype == np.uint16:
return _UINT16
if dtype == np.uint32:
return _UINT32
if dtype == np.uint64:
return _UINT64
if dtype == np.float32:
return _FLOAT32
if dtype == np.float64:
return _FLOAT64
msg = f"unsupported dtype: {dtype}"
raise TypeError(msg)
def cdist(
queries,
choices,
*,
scorer=ratio,
processor=None,
score_cutoff=None,
score_hint=None,
score_multiplier=1,
dtype=None,
workers=1,
**kwargs,
):
import numpy as np
dtype = _dtype_to_type_num(dtype)
return np.asarray(
_cdist(
queries,
choices,
scorer=scorer,
processor=processor,
score_cutoff=score_cutoff,
score_hint=score_hint,
score_multiplier=score_multiplier,
dtype=dtype,
workers=workers,
**kwargs,
)
)
cdist.__doc__ = _cdist.__doc__
def cpdist(
queries,
choices,
*,
scorer=ratio,
processor=None,
score_cutoff=None,
score_hint=None,
score_multiplier=1,
dtype=None,
workers=1,
**kwargs,
):
import numpy as np
dtype = _dtype_to_type_num(dtype)
distance_matrix = _cpdist(
queries,
choices,
scorer=scorer,
processor=processor,
score_cutoff=score_cutoff,
score_hint=score_hint,
score_multiplier=score_multiplier,
dtype=dtype,
workers=workers,
**kwargs,
)
return np.asarray(distance_matrix)
cpdist.__doc__ = _cpdist.__doc__