From 2539bc8f573a1233e5a0b25ad83d26caaa8eaba6 Mon Sep 17 00:00:00 2001 From: David Manouchehri Date: Wed, 9 Jul 2025 15:17:31 +0000 Subject: [PATCH 1/2] Enhance airport search to support multiple search criteria - Add search by airport code (e.g., GCM) - case-sensitive - Add search by ICAO code (e.g., MWCR) - case-sensitive - Add search by country code (e.g., KY) - case-sensitive - Add search by city name (e.g., George Town) - case-insensitive - Maintain backward compatibility with enum name search - Add comprehensive test file demonstrating all search methods --- fast_flights/airport_data.py | 97 +++++++++++++++++++++++++++ fast_flights/search.py | 53 +++++++++++++-- test_airport_search.py | 123 +++++++++++++++++++++++++++++++++++ 3 files changed, 266 insertions(+), 7 deletions(-) create mode 100644 fast_flights/airport_data.py create mode 100644 test_airport_search.py diff --git a/fast_flights/airport_data.py b/fast_flights/airport_data.py new file mode 100644 index 00000000..db51ad7f --- /dev/null +++ b/fast_flights/airport_data.py @@ -0,0 +1,97 @@ +import csv +import os +from dataclasses import dataclass +from typing import Dict, List, Optional +from ._generated_enum import Airport + + +@dataclass +class AirportInfo: + code: str + name: str + city: Optional[str] + country_id: str + icao: Optional[str] + enum_member: Optional[Airport] + + +class AirportData: + def __init__(self): + self.airports: List[AirportInfo] = [] + self.by_code: Dict[str, AirportInfo] = {} + self.by_icao: Dict[str, AirportInfo] = {} + self.by_country: Dict[str, List[AirportInfo]] = {} + self.by_city_lower: Dict[str, List[AirportInfo]] = {} + self._load_data() + + def _load_data(self): + csv_path = os.path.join(os.path.dirname(__file__), '..', 'enums', 'airports.csv') + + with open(csv_path, 'r', encoding='utf-8') as f: + reader = csv.DictReader(f) + for row in reader: + code = row['code'] + name = row['name'] + city = row.get('city', '').strip() or None + country_id = row['country_id'] + icao = row.get('icao', '').strip() or None + + # Try to find corresponding enum member + enum_member = None + if 'AIRPORT' in name.upper(): + normalized_name = '_'.join( + name.replace('-', ' ') + .replace('.', ' ') + .replace('/', ' ') + .replace("'", '') + .replace('(', ' ') + .replace(')', ' ') + .replace('–', ' ') + .split() + ).upper() + + try: + enum_member = getattr(Airport, normalized_name, None) + except AttributeError: + pass + + airport_info = AirportInfo( + code=code, + name=name, + city=city, + country_id=country_id, + icao=icao, + enum_member=enum_member + ) + + self.airports.append(airport_info) + + # Index by code + self.by_code[code] = airport_info + + # Index by ICAO + if icao: + self.by_icao[icao] = airport_info + + # Index by country + if country_id not in self.by_country: + self.by_country[country_id] = [] + self.by_country[country_id].append(airport_info) + + # Index by city (lowercase for case-insensitive search) + if city: + city_lower = city.lower() + if city_lower not in self.by_city_lower: + self.by_city_lower[city_lower] = [] + self.by_city_lower[city_lower].append(airport_info) + + +# Singleton instance +_airport_data: Optional[AirportData] = None + + +def get_airport_data() -> AirportData: + global _airport_data + if _airport_data is None: + _airport_data = AirportData() + return _airport_data \ No newline at end of file diff --git a/fast_flights/search.py b/fast_flights/search.py index 87fb0315..c777154a 100644 --- a/fast_flights/search.py +++ b/fast_flights/search.py @@ -1,18 +1,57 @@ -from typing import List +from typing import List, Set from ._generated_enum import Airport +from .airport_data import get_airport_data def search_airport(query: str) -> List[Airport]: """Search for airports. Args: - query (str): The query. + query (str): The query. Can match: + - Airport enum member names (case-insensitive) + - Airport codes like GCM (case-sensitive) + - ICAO codes like MWCR (case-sensitive) + - Country codes like KY (case-sensitive) + - City names like George Town (case-insensitive) Returns: list[Airport]: A list of airports (enum `Airports`). """ - return [ - ref - for aname, ref in Airport.__members__.items() - if query.lower() in aname.lower() - ] + results: Set[Airport] = set() + airport_data = get_airport_data() + + # Search by enum member name (case-insensitive) - original behavior + # Also handle spaces as underscores for better matching + query_normalized = query.lower().replace(' ', '_') + for aname, ref in Airport.__members__.items(): + aname_lower = aname.lower() + if query.lower() in aname_lower or query_normalized in aname_lower: + results.add(ref) + + # Search by airport code (case-sensitive) + if query in airport_data.by_code: + airport_info = airport_data.by_code[query] + if airport_info.enum_member: + results.add(airport_info.enum_member) + + # Search by ICAO code (case-sensitive) + if query in airport_data.by_icao: + airport_info = airport_data.by_icao[query] + if airport_info.enum_member: + results.add(airport_info.enum_member) + + # Search by country code (case-sensitive) + if query in airport_data.by_country: + for airport_info in airport_data.by_country[query]: + if airport_info.enum_member: + results.add(airport_info.enum_member) + + # Search by city name (case-insensitive) + query_lower = query.lower() + for city_lower, airports in airport_data.by_city_lower.items(): + if query_lower in city_lower: + for airport_info in airports: + if airport_info.enum_member: + results.add(airport_info.enum_member) + + return list(results) diff --git a/test_airport_search.py b/test_airport_search.py new file mode 100644 index 00000000..631d19bd --- /dev/null +++ b/test_airport_search.py @@ -0,0 +1,123 @@ +#!/usr/bin/env python3 +""" +Test file for enhanced airport search functionality in fast-flights +Uses Owen Roberts International Airport (GCM) in Cayman Islands as the main example +""" + +from fast_flights import search_airport + +def print_results(query, results, max_display=5): + """Helper function to print search results""" + print(f"\nSearch for '{query}': {len(results)} result(s)") + for i, airport in enumerate(results[:max_display]): + print(f" {i+1}. {airport}") + if len(results) > max_display: + print(f" ... and {len(results) - max_display} more") + +def main(): + print("Testing enhanced airport search functionality") + print("=" * 80) + print("Main example: Owen Roberts International Airport (GCM) in George Town, Cayman Islands") + print("=" * 80) + + # Test 1: Search by airport code (case-sensitive) + print("\n1. SEARCH BY AIRPORT CODE (case-sensitive)") + print("-" * 40) + results = search_airport("GCM") + print_results("GCM", results) + + # Test case sensitivity - should not match + results = search_airport("gcm") + print_results("gcm", results) + + # Test 2: Search by ICAO code (case-sensitive) + print("\n\n2. SEARCH BY ICAO CODE (case-sensitive)") + print("-" * 40) + results = search_airport("MWCR") + print_results("MWCR", results) + + # Test case sensitivity - should not match + results = search_airport("mwcr") + print_results("mwcr", results) + + # Test 3: Search by country code (case-sensitive) + print("\n\n3. SEARCH BY COUNTRY CODE (case-sensitive)") + print("-" * 40) + results = search_airport("KY") + print_results("KY", results) + print("\nNote: KY is the country code for Cayman Islands (3 airports)") + print("Additional results come from airports with 'KY' in their enum names (e.g., KYIV, KYZYL)") + + # Test case sensitivity - should not match country code but might match in enum names + results = search_airport("ky") + print_results("ky", results) + print("Note: Lowercase 'ky' doesn't match country code but matches enum names") + + # Test 4: Search by city name (case-insensitive) + print("\n\n4. SEARCH BY CITY NAME (case-insensitive)") + print("-" * 40) + results = search_airport("George Town") + print_results("George Town", results) + print("\nNote: George Town exists in multiple countries (Cayman Islands, Bahamas, etc.)") + + # Test case insensitivity + results = search_airport("george town") + print_results("george town", results) + + # Test partial match + results = search_airport("george") + print_results("george", results, max_display=10) + + # Test 5: Search by enum member name (case-insensitive) + print("\n\n5. SEARCH BY ENUM MEMBER NAME (case-insensitive)") + print("-" * 40) + results = search_airport("owen roberts") + print_results("owen roberts", results) + + results = search_airport("OWEN ROBERTS") + print_results("OWEN ROBERTS", results) + + results = search_airport("owen_roberts") + print_results("owen_roberts", results) + + # Test 6: Edge cases + print("\n\n6. EDGE CASES") + print("-" * 40) + + # Non-existent search + results = search_airport("ZZZZZ") + print_results("ZZZZZ", results) + + # Very common word that might match many airports + results = search_airport("international") + print(f"\nSearch for 'international': {len(results)} result(s)") + print("(Too many to display - this matches all airports with 'International' in their name)") + + # Empty string + results = search_airport("") + print(f"\nSearch for '' (empty string): {len(results)} result(s)") + + # Test 7: Other Cayman Islands airports + print("\n\n7. OTHER CAYMAN ISLANDS AIRPORTS") + print("-" * 40) + print("Getting actual Cayman Islands airports only (not just 'KY' matches)...") + + # Import to access the data directly + from fast_flights.airport_data import get_airport_data + airport_data = get_airport_data() + cayman_airports = airport_data.by_country.get("KY", []) + + print(f"\nFound {len(cayman_airports)} airport(s) in Cayman Islands:") + for i, airport_info in enumerate(cayman_airports, 1): + print(f" {i}. {airport_info.enum_member}") + print(f" - {airport_info.name} ({airport_info.code})") + if airport_info.icao: + print(f" - ICAO: {airport_info.icao}") + if airport_info.city: + print(f" - City: {airport_info.city}") + + print("\n" + "=" * 80) + print("Test completed successfully!") + +if __name__ == "__main__": + main() \ No newline at end of file From 68aaaf3d11123b1773fb2cc0ccdeba97e05c06b7 Mon Sep 17 00:00:00 2001 From: David Manouchehri Date: Wed, 9 Jul 2025 18:30:59 +0000 Subject: [PATCH 2/2] Fix missing airports.csv in PyPI package distribution Move airports.csv to fast_flights package directory and configure pyproject.toml to include it as package data for proper distribution. --- fast_flights/airport_data.py | 2 +- {enums => fast_flights}/airports.csv | 0 pyproject.toml | 6 ++++++ 3 files changed, 7 insertions(+), 1 deletion(-) rename {enums => fast_flights}/airports.csv (100%) diff --git a/fast_flights/airport_data.py b/fast_flights/airport_data.py index db51ad7f..94a27475 100644 --- a/fast_flights/airport_data.py +++ b/fast_flights/airport_data.py @@ -25,7 +25,7 @@ def __init__(self): self._load_data() def _load_data(self): - csv_path = os.path.join(os.path.dirname(__file__), '..', 'enums', 'airports.csv') + csv_path = os.path.join(os.path.dirname(__file__), 'airports.csv') with open(csv_path, 'r', encoding='utf-8') as f: reader = csv.DictReader(f) diff --git a/enums/airports.csv b/fast_flights/airports.csv similarity index 100% rename from enums/airports.csv rename to fast_flights/airports.csv diff --git a/pyproject.toml b/pyproject.toml index ab85e376..581de6fb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,6 +39,12 @@ packages = [ "fast_flights" ] +[tool.flit.module] +name = "fast_flights" + +[tool.flit.sdist] +include = ["fast_flights/airports.csv"] + [tool.pyright] pythonVersion = '3.8'