diff --git a/outlines/types/airports.py b/outlines/types/airports.py index b4ed6784d..677ec55db 100644 --- a/outlines/types/airports.py +++ b/outlines/types/airports.py @@ -1,10 +1,9 @@ """Generate valid airport codes.""" from enum import Enum -from pyairports.airports import AIRPORT_LIST +import airportsdata + +AIRPORT_IATA_LIST = [(iata, iata) for iata in sorted(airportsdata.load().keys())] -AIRPORT_IATA_LIST = list( - {(airport[3], airport[3]) for airport in AIRPORT_LIST if airport[3] != ""} -) IATA = Enum("Airport", AIRPORT_IATA_LIST) # type:ignore diff --git a/pyproject.toml b/pyproject.toml index dfda81acd..4d87e1fca 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,7 +40,7 @@ dependencies = [ "datasets", "typing_extensions", "pycountry", - "pyairports", + "airportsdata", ] dynamic = ["version"] @@ -137,7 +137,7 @@ module = [ "uvicorn.*", "fastapi.*", "pycountry.*", - "pyairports.*", + "airportsdata.*", ] ignore_missing_imports = true