Enhance airport search to support multiple search criteria

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__), '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

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)

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'
 

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()