avwx.station.search

 1"""Station text-based search."""
 2
 3# stdlib
 4from __future__ import annotations
 5
 6from contextlib import suppress
 7from functools import lru_cache
 8from typing import TYPE_CHECKING
 9
10# module
11from avwx.exceptions import MissingExtraModule
12from avwx.load_utils import LazyCalc
13from avwx.station.meta import STATIONS
14from avwx.station.station import Station, station_filter
15
16if TYPE_CHECKING:
17    from collections.abc import Iterable
18
19# Catch import error only if user attemps a text search
20with suppress(ModuleNotFoundError):
21    from rapidfuzz import fuzz, process, utils
22
23
24TYPE_ORDER = [
25    "large_airport",
26    "medium_airport",
27    "small_airport",
28    "seaplane_base",
29    "heliport",
30    "balloonport",
31    "weather_station",
32]
33
34
35def _format_search(airport: dict, keys: Iterable[str]) -> str | None:
36    values = [airport.get(k) for k in keys]
37    code = values[0] or values[2]
38    if not code:
39        return None
40    values.insert(0, code)
41    return " - ".join(k for k in values if k)
42
43
44def _build_corpus() -> list[str]:
45    keys = ("icao", "iata", "gps", "local", "city", "state", "name")
46    return [text for s in STATIONS.values() if (text := _format_search(s, keys))]
47
48
49_CORPUS = LazyCalc(_build_corpus)
50
51
52def _sort_key(result: tuple[Station, int]) -> tuple[int, ...]:
53    station, score = result
54    try:
55        type_order = TYPE_ORDER.index(station.type)
56    except ValueError:
57        type_order = 10
58    return (score, 10 - type_order)
59
60
61@lru_cache(maxsize=128)
62def search(
63    text: str,
64    limit: int = 10,
65    *,
66    is_airport: bool = False,
67    sends_reports: bool = True,
68) -> list[Station]:
69    """Text search for stations against codes, name, city, and state.
70
71    Results may be shorter than limit value.
72    """
73    try:
74        results = process.extract(
75            text,
76            _CORPUS.value,
77            limit=limit * 20,
78            scorer=fuzz.token_set_ratio,
79            processor=utils.default_process,
80        )
81    except NameError as name_error:
82        extra = "fuzz"
83        raise MissingExtraModule(extra) from name_error
84    results = [(Station.from_code(k[:4]), s) for k, s, _ in results]
85    results.sort(key=_sort_key, reverse=True)
86    results = [s for s, _ in results if station_filter(s, is_airport=is_airport, reporting=sends_reports)]
87    return results[:limit] if len(results) > limit else results