"""Module contains the async interface to match needle against haystack in batch."""
import asyncio
import heapq
from typing import Any, Callable, Dict, List, Union, cast
from pfzy.score import SCORE_INDICES, fzy_scorer
from pfzy.types import HAYSTACKS
async def _rank_task(
scorer: Callable[[str, str], SCORE_INDICES],
needle: str,
haystacks: List[Union[str, Dict[str, Any]]],
key: str,
) -> List[Dict[str, Any]]:
"""Calculate the score for needle against given list of haystacks and rank them.
Args:
scorer: Scorer to be used to do the calculation.
needle: Substring to search.
haystacks: List of dictionary containing the haystack to be searched.
key: The key within the `haystacks` dictionary that contains the actual string value.
Return:
Sorted list of haystacks based on the needle score with additional keys for the `score`
and `indices`.
"""
result = []
for haystack in haystacks:
score, indices = scorer(needle, cast(Dict, haystack)[key])
if indices is None:
continue
result.append(
{
"score": score,
"indices": indices,
"haystack": haystack,
}
)
result.sort(key=lambda x: x["score"], reverse=True)
return result
[docs]async def fuzzy_match(
needle: str,
haystacks: HAYSTACKS,
key: str = "",
batch_size: int = 4096,
scorer: Callable[[str, str], SCORE_INDICES] = None,
) -> List[Dict[str, Any]]:
"""Fuzzy find the needle within list of haystacks and get matched results with matching index.
Note:
The `key` argument is optional when the provided `haystacks` argument is a list of :class:`str`.
It will be given a default key `value` if not present.
Warning:
The `key` argument is required when provided `haystacks` argument is a list of :class:`dict`.
If not present, :class:`TypeError` will be raised.
Args:
needle: String to search within the `haystacks`.
haystacks: List of haystack/longer strings to be searched.
key: If `haystacks` is a list of dictionary, provide the key that
can obtain the haystack value to search.
batch_size: Number of entry to be processed together.
scorer (Callable[[str, str], SCORE_indices]): Desired scorer to use. Currently only :func:`~pfzy.score.fzy_scorer` and :func:`~pfzy.score.substr_scorer` is supported.
Raises:
TypeError: When the argument `haystacks` is :class:`list` of :class:`dict` and the `key` argument
is missing, :class:`TypeError` will be raised.
Returns:
List of matching `haystacks` with additional key indices and score.
Examples:
>>> import asyncio
>>> asyncio.run(fuzzy_match("ab", ["acb", "acbabc"]))
[{'value': 'acbabc', 'indices': [3, 4]}, {'value': 'acb', 'indices': [0, 2]}]
"""
if scorer is None:
scorer = fzy_scorer
for index, haystack in enumerate(haystacks):
if not isinstance(haystack, dict):
if not key:
key = "value"
haystacks[index] = {key: haystack}
if not key:
raise TypeError(
f"${fuzzy_match.__name__} missing 1 required argument: 'key', 'key' is required when haystacks is an instance of dict"
)
batches = await asyncio.gather(
*(
_rank_task(
scorer,
needle,
haystacks[offset : offset + batch_size],
key,
)
for offset in range(0, len(haystacks), batch_size)
)
)
results = heapq.merge(*batches, key=lambda x: x["score"], reverse=True)
choices = []
for candidate in results:
candidate["haystack"]["indices"] = candidate["indices"]
choices.append(candidate["haystack"])
return choices