avwx.station.search

Station text-based search

View Source

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

TYPE_ORDER = ['large_airport', 'medium_airport', 'small_airport', 'seaplane_base', 'heliport', 'balloonport', 'weather_station']

@lru_cache(maxsize=128)

def search( text: str, limit: int = 10, is_airport: bool = False, sends_reports: bool = True) -> List[avwx.station.station.Station]: View Source

59@lru_cache(maxsize=128)
60def search(
61    text: str,
62    limit: int = 10,
63    is_airport: bool = False,
64    sends_reports: bool = True,
65) -> List[Station]:
66    """Text search for stations against codes, name, city, and state
67
68    Results may be shorter than limit value
69    """
70    try:
71        results = process.extract(
72            text,
73            _CORPUS.value,
74            limit=limit * 20,
75            scorer=fuzz.token_set_ratio,
76            processor=utils.default_process,
77        )
78    except NameError as name_error:
79        raise ModuleNotFoundError(
80            'rapidfuzz must be installed to use text search. Run "pip install avwx-engine[fuzz]" to enable this'
81        ) from name_error
82    results = [(Station.from_code(k[:4]), s) for k, s, _ in results]
83    results.sort(key=_sort_key, reverse=True)
84    results = [s for s, _ in results if station_filter(s, is_airport, sends_reports)]
85    return results[:limit] if len(results) > limit else results

Text search for stations against codes, name, city, and state

Results may be shorter than limit value