| """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 |
|
|
|
|
| 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 |
|
|
