flight_search.py

import os
import requests
from dotenv import load_dotenv
from flight_search_token import GetToken
from sheety_data_manager import SheetyDataManager

# Load environment variables from .env file
load_dotenv()

# API endpoint for city search
city_search_url = "https://test.api.amadeus.com/v1/reference-data/locations/cities"


class FlightSearch:
    def __init__(self, token: GetToken, data: SheetyDataManager):
        """
        Initialize FlightSearch with token and data manager.
        :param token: Object to retrieve API token
        :param data: Object to interact with Sheety data manager
        """
        self.token = token.returned_token()
        self.data = data

    def test(self, city):
        """
        Test function for debugging purposes.
        :param city: City name
        :return: A simple test message
        """
        return "TESTING"

    def search_iata(self, city_name):
        """
        Search for the IATA code of a given city using the Amadeus API.
        :param city_name: Name of the city to search
        :return: IATA code of the city if found, None otherwise
        """
        parameters = {
            "keyword": city_name,  # City name to search for
            "include": "AIRPORTS",  # Include airport details in the results
            "max": 2  # Maximum number of results
        }

        headers = {
            "Authorization": f"Bearer {self.token}"  # Authorization header with the API token
        }

        try:
            # Make a GET request to the Amadeus API
            response = requests.get(url=city_search_url, headers=headers, params=parameters)
            response.raise_for_status()  # Raise an error for unsuccessful requests
            data = response.json()["data"]  # Extract the data from the response

            # Extract the IATA code from the first valid result
            for code in data:
                try:
                    iata_code = code["iataCode"]
                    return iata_code
                except KeyError:
                    print("IATA code not found in the response.")
        except requests.RequestException as e:
            print(f"API Request failed: {e}")
        return None


# Example usage of the FlightSearch class
if __name__ == "__main__":
    # Initialize token and data manager
    token_ = GetToken()
    sheet = SheetyDataManager()

    # Create an instance of the FlightSearch class
    flight_search = FlightSearch(token=token_, data=sheet)

    # Search for the IATA code of a city
    iata_code = flight_search.search_iata("paris")
    print(iata_code)