park-search/pota_search.py
2022-12-28 16:16:31 -08:00

92 lines
2.2 KiB
Python

import csv
import sys
import dataclasses
import thefuzz.fuzz, thefuzz.process
import typing
import itertools
@dataclasses.dataclass
class Park:
reference: str
name: str
location: typing.List[str]
latitude: float
longitude: float
grid: str
attempts: int
activations: int
qsos: int
@dataclasses.dataclass
class ParkDatabase:
parks: typing.List[Park]
def __post_init__(self):
self.parks_dict = {park.reference: park for park in self.parks}
def search_name(self, string):
return [
self.parks_dict[i[0].split("|")[0]]
for i in thefuzz.process.extract(
string,
[park.reference + "|" + park.name for park in self.parks],
scorer=thefuzz.fuzz.partial_ratio,
)
]
def filter_location(self, location):
location = location.upper()
return ParkDatabase(
[
park
for park in self.parks
if (location in park.location)
or (
location
in [
park_location.split("-")[0]
for park_location in park.location
]
)
]
)
def get_locations(self):
return _uniq(itertools.chain(*[park.location for park in self.parks]))
def get_expanded_locations(self):
return _uniq(_expand_locations(self.get_locations()))
def _uniq(data):
return sorted(set(data))
def _expand_location(code):
return code, code.split("-")[0]
def _expand_locations(codes):
return itertools.chain(*[_expand_location(code) for code in codes])
def load_csv(f):
return ParkDatabase(
[
Park(
park["reference"],
park["name"],
park["locationDesc"].split(","),
float(park["latitude"]),
float(park["longitude"]),
park["grid"],
int(park["attempts"]),
int(park["activations"]),
int(park["qsos"]),
)
for park in csv.DictReader(f)
]
)