From ac18d60f8c78a0530e01e0728831915d68980b26 Mon Sep 17 00:00:00 2001 From: x Date: Wed, 11 Sep 2024 16:53:03 -0400 Subject: [PATCH] added max_stops --- .gitignore | 1 + README.md | 7 ++- example.py | 83 ++++++++++++++++++++++++++++++++++++ fast_flights/core.py | 17 +++++--- fast_flights/filter.py | 8 +++- fast_flights/flights.proto | 7 +-- fast_flights/flights_impl.py | 31 +++++++++++--- fast_flights/flights_pb2.py | 38 ++++++++--------- 8 files changed, 155 insertions(+), 37 deletions(-) create mode 100644 example.py diff --git a/.gitignore b/.gitignore index 9f161b2e..0fd49f58 100644 --- a/.gitignore +++ b/.gitignore @@ -1,3 +1,4 @@ **/__pycache__/ /main.py .ruff/ +.idea \ No newline at end of file diff --git a/README.md b/README.md index 2b52d6f7..598c25b0 100644 --- a/README.md +++ b/README.md @@ -16,7 +16,7 @@ $ pip install fast-flights ## Basic To use `fast-flights`, you'll first create a filter (inherited from `?tfs=`) to perform a request. -Then, add `flight_data`, `trip`, `seat` and `passengers` info to use the API directly. +Then, add `flight_data`, `trip`, `seat`, `passengers`, and (optional) `max_stops` info to use the API directly. Honorable mention: I like birds. Yes, I like birds. @@ -42,6 +42,7 @@ filter = create_filter( infants_in_seat=0, infants_on_lap=0 ), + max_stops=args.max_stops ) # Get flights with a filter @@ -51,6 +52,10 @@ result = get_flights(filter) print("The price is currently", result.current_price) ``` +A command-line example script is included as `example.py`. Usage is as follows: + +`python3 example.py --origin LAX --destination LGA --depart_date 2025-2-26 --return_date 2025-02-29 --max_stops 1` + **Information**: Display additional information. ```python # Get the first flight diff --git a/example.py b/example.py new file mode 100644 index 00000000..9721b567 --- /dev/null +++ b/example.py @@ -0,0 +1,83 @@ +import argparse +import json +from fast_flights import FlightData, Passengers, create_filter, get_flights + +def flight_to_dict(flight): + return { + "is_best": getattr(flight, 'is_best', None), + "name": getattr(flight, 'name', None), + "departure": getattr(flight, 'departure', None), + "arrival": getattr(flight, 'arrival', None), + "arrival_time_ahead": getattr(flight, 'arrival_time_ahead', None), + "duration": getattr(flight, 'duration', None), + "stops": getattr(flight, 'stops', None), + "delay": getattr(flight, 'delay', None), + "price": getattr(flight, 'price', None), + } + +def result_to_dict(result): + return { + "current_price": getattr(result, 'current_price', None), + "flights": [flight_to_dict(flight) for flight in getattr(result, 'flights', [])] + } + +def main(): + # Argument parser for command-line input + parser = argparse.ArgumentParser(description="Flight Price Finder") + parser.add_argument('--origin', required=True, help="Origin airport code") + parser.add_argument('--destination', required=True, help="Destination airport code") + parser.add_argument('--depart_date', required=True, help="Beginning trip date (YYYY-MM-DD)") + parser.add_argument('--return_date', required=True, help="Ending trip date (YYYY-MM-DD)") + parser.add_argument('--adults', type=int, default=1, help="Number of adult passengers") + parser.add_argument('--type', type=str, default="economy", help="Fare class (economy, premium-economy, business or first)") + parser.add_argument('--max_stops', type=int, help="Maximum number of stops (optional, [0|1|2])") + + args = parser.parse_args() + + # Create a new filter + filter = create_filter( + flight_data=[ + FlightData( + date=args.depart_date, # Date of departure for outbound flight + from_airport=args.origin, + to_airport=args.destination + ), + FlightData( + date=args.return_date, # Date of departure for return flight + from_airport=args.destination, + to_airport=args.origin + ), + ], + trip="round-trip", # Trip (round-trip, one-way) + seat=args.type, # Seat (economy, premium-economy, business or first) + passengers=Passengers( + adults=args.adults, + children=0, + infants_in_seat=0, + infants_on_lap=0 + ), + max_stops=args.max_stops + ) + + b64 = filter.as_b64().decode('utf-8') + print( + "https://www.google.com/travel/flights?tfs=%s" % b64 + ) + + # Get flights with the filter + result = get_flights(filter) + + try: + # Manually convert the result to a dictionary before serialization + result_dict = result_to_dict(result) + print(json.dumps(result_dict, indent=4)) + except TypeError as e: + print("Serialization to JSON failed. Raw result output:") + print(result) + print("Error details:", str(e)) + + # Print price information + print("The price is currently", result.current_price) + +if __name__ == "__main__": + main() diff --git a/fast_flights/core.py b/fast_flights/core.py index 1e0a9924..01697c16 100644 --- a/fast_flights/core.py +++ b/fast_flights/core.py @@ -16,18 +16,23 @@ def request_flights( tfs: TFSData, *, + max_stops: Optional[int] = None, # Optionally pass max_stops if needed currency: Optional[str] = None, language: Optional[str], **kwargs: Any, ) -> requests.Response: + params = { + "tfs": tfs.as_b64(), + "hl": language, + "tfu": "EgQIABABIgA", # show all flights and prices condition + "curr": currency, + } + if max_stops is not None: + params["max_stops"] = max_stops # Handle max_stops in request + r = requests.get( "https://www.google.com/travel/flights", - params={ - "tfs": tfs.as_b64(), - "hl": language, - "tfu": "EgQIABABIgA", # show all flights and prices condition - "curr": currency, - }, + params=params, headers={"user-agent": ua, "accept-language": "en"}, **kwargs, ) diff --git a/fast_flights/filter.py b/fast_flights/filter.py index 0b2e9f62..503db3c6 100644 --- a/fast_flights/filter.py +++ b/fast_flights/filter.py @@ -1,13 +1,13 @@ -from typing import Literal, List +from typing import Literal, List, Optional from .flights_impl import FlightData, Passengers, TFSData - def create_filter( *, flight_data: List[FlightData], trip: Literal["round-trip", "one-way", "multi-city"], passengers: Passengers, seat: Literal["economy", "premium-economy", "business", "first"], + max_stops: Optional[int] = None, # New argument ) -> TFSData: """Create a filter. (``?tfs=``) @@ -16,7 +16,11 @@ def create_filter( trip ("one-way" | "round-trip" | "multi-city"): Trip type. passengers (Passengers): Passengers. seat ("economy" | "premium-economy" | "business" | "first"): Seat. + max_stops (int, optional): Maximum number of stops. Defaults to None. """ + for fd in flight_data: + fd.max_stops = max_stops + return TFSData.from_interface( flight_data=flight_data, trip=trip, passengers=passengers, seat=seat ) diff --git a/fast_flights/flights.proto b/fast_flights/flights.proto index ff225f35..b65a93ac 100644 --- a/fast_flights/flights.proto +++ b/fast_flights/flights.proto @@ -5,9 +5,10 @@ message Airport { } message FlightData { - string date = 2; - Airport from_flight = 13; - Airport to_flight = 14; + string date = 2; + Airport from_flight = 13; + Airport to_flight = 14; + optional int32 max_stops = 5; } enum Seat { diff --git a/fast_flights/flights_impl.py b/fast_flights/flights_impl.py index a58030fe..ef7e84f5 100644 --- a/fast_flights/flights_impl.py +++ b/fast_flights/flights_impl.py @@ -1,7 +1,7 @@ """Typed implementation of flights_pb2.py""" import base64 -from typing import Any, List, TYPE_CHECKING, Literal, Union +from typing import Any, List, Optional, TYPE_CHECKING, Literal, Union from . import flights_pb2 as PB from ._generated_enum import Airport @@ -17,12 +17,14 @@ class FlightData: date (str): Date. from_airport (str): Departure (airport). Where from? to_airport (str): Arrival (airport). Where to? + max_stops (int, optional): Maximum number of stops. Default is None. """ - __slots__ = ("date", "from_airport", "to_airport") + __slots__ = ("date", "from_airport", "to_airport", "max_stops") date: str from_airport: str to_airport: str + max_stops: Optional[int] def __init__( self, @@ -30,6 +32,7 @@ def __init__( date: str, from_airport: Union[Airport, str], to_airport: Union[Airport, str], + max_stops: Optional[int] = None, ): self.date = date self.from_airport = ( @@ -38,18 +41,22 @@ def __init__( self.to_airport = ( to_airport.value if isinstance(to_airport, Airport) else to_airport ) + self.max_stops = max_stops def attach(self, info: PB.Info) -> None: # type: ignore data = info.data.add() data.date = self.date data.from_flight.airport = self.from_airport data.to_flight.airport = self.to_airport + if self.max_stops is not None: + data.max_stops = self.max_stops def __repr__(self) -> str: return ( f"FlightData(date={self.date!r}, " f"from_airport={self.from_airport}, " - f"to_airport={self.to_airport})" + f"to_airport={self.to_airport}, " + f"max_stops={self.max_stops})" ) @@ -98,11 +105,13 @@ def __init__( seat: PB.Seat, # type: ignore trip: PB.Trip, # type: ignore passengers: Passengers, + max_stops: Optional[int] = None, # Add max_stops to the constructor ): self.flight_data = flight_data self.seat = seat self.trip = trip self.passengers = passengers + self.max_stops = max_stops # Store max_stops def pb(self) -> PB.Info: # type: ignore info = PB.Info() @@ -114,6 +123,11 @@ def pb(self) -> PB.Info: # type: ignore for fd in self.flight_data: fd.attach(info) + # If max_stops is set, attach it to all flight data entries + if self.max_stops is not None: + for flight in info.data: + flight.max_stops = self.max_stops + return info def to_string(self) -> bytes: @@ -129,6 +143,7 @@ def from_interface( trip: Literal["round-trip", "one-way", "multi-city"], passengers: Passengers, seat: Literal["economy", "premium-economy", "business", "first"], + max_stops: Optional[int] = None, # Add max_stops to the method signature ): """Use ``?tfs=`` from an interface. @@ -137,6 +152,7 @@ def from_interface( trip ("one-way" | "round-trip" | "multi-city"): Trip type. passengers (Passengers): Passengers. seat ("economy" | "premium-economy" | "business" | "first"): Seat. + max_stops (int, optional): Maximum number of stops. """ trip_t = { "round-trip": PB.Trip.ROUND_TRIP, @@ -151,8 +167,13 @@ def from_interface( }[seat] return TFSData( - flight_data=flight_data, seat=seat_t, trip=trip_t, passengers=passengers + flight_data=flight_data, + seat=seat_t, + trip=trip_t, + passengers=passengers, + max_stops=max_stops # Pass max_stops into TFSData ) def __repr__(self) -> str: - return f"TFSData({'hello'!r}, flight_data={self.flight_data!r})" + return f"TFSData(flight_data={self.flight_data!r}, max_stops={self.max_stops!r})" + diff --git a/fast_flights/flights_pb2.py b/fast_flights/flights_pb2.py index f1bca6b9..ff7fd3ef 100644 --- a/fast_flights/flights_pb2.py +++ b/fast_flights/flights_pb2.py @@ -1,10 +1,7 @@ -# type: ignore - # -*- coding: utf-8 -*- # Generated by the protocol buffer compiler. DO NOT EDIT! # source: flights.proto """Generated protocol buffer code.""" - from google.protobuf.internal import builder as _builder from google.protobuf import descriptor as _descriptor from google.protobuf import descriptor_pool as _descriptor_pool @@ -14,24 +11,25 @@ _sym_db = _symbol_database.Default() -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile( - b'\n\rflights.proto"\x1a\n\x07\x41irport\x12\x0f\n\x07\x61irport\x18\x02 \x01(\t"V\n\nFlightData\x12\x0c\n\x04\x64\x61te\x18\x02 \x01(\t\x12\x1d\n\x0b\x66rom_flight\x18\r \x01(\x0b\x32\x08.Airport\x12\x1b\n\tto_flight\x18\x0e \x01(\x0b\x32\x08.Airport"k\n\x04Info\x12\x19\n\x04\x64\x61ta\x18\x03 \x03(\x0b\x32\x0b.FlightData\x12\x13\n\x04seat\x18\t \x01(\x0e\x32\x05.Seat\x12\x1e\n\npassengers\x18\x08 \x03(\x0e\x32\n.Passenger\x12\x13\n\x04trip\x18\x13 \x01(\x0e\x32\x05.Trip*S\n\x04Seat\x12\x10\n\x0cUNKNOWN_SEAT\x10\x00\x12\x0b\n\x07\x45\x43ONOMY\x10\x01\x12\x13\n\x0fPREMIUM_ECONOMY\x10\x02\x12\x0c\n\x08\x42USINESS\x10\x03\x12\t\n\x05\x46IRST\x10\x04*E\n\x04Trip\x12\x10\n\x0cUNKNOWN_TRIP\x10\x00\x12\x0e\n\nROUND_TRIP\x10\x01\x12\x0b\n\x07ONE_WAY\x10\x02\x12\x0e\n\nMULTI_CITY\x10\x03*_\n\tPassenger\x12\x15\n\x11UNKNOWN_PASSENGER\x10\x00\x12\t\n\x05\x41\x44ULT\x10\x01\x12\t\n\x05\x43HILD\x10\x02\x12\x12\n\x0eINFANT_IN_SEAT\x10\x03\x12\x11\n\rINFANT_ON_LAP\x10\x04\x62\x06proto3' -) + + +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\rflights.proto\"\x1a\n\x07\x41irport\x12\x0f\n\x07\x61irport\x18\x02 \x01(\t\"|\n\nFlightData\x12\x0c\n\x04\x64\x61te\x18\x02 \x01(\t\x12\x1d\n\x0b\x66rom_flight\x18\r \x01(\x0b\x32\x08.Airport\x12\x1b\n\tto_flight\x18\x0e \x01(\x0b\x32\x08.Airport\x12\x16\n\tmax_stops\x18\x05 \x01(\x05H\x00\x88\x01\x01\x42\x0c\n\n_max_stops\"k\n\x04Info\x12\x19\n\x04\x64\x61ta\x18\x03 \x03(\x0b\x32\x0b.FlightData\x12\x13\n\x04seat\x18\t \x01(\x0e\x32\x05.Seat\x12\x1e\n\npassengers\x18\x08 \x03(\x0e\x32\n.Passenger\x12\x13\n\x04trip\x18\x13 \x01(\x0e\x32\x05.Trip*S\n\x04Seat\x12\x10\n\x0cUNKNOWN_SEAT\x10\x00\x12\x0b\n\x07\x45\x43ONOMY\x10\x01\x12\x13\n\x0fPREMIUM_ECONOMY\x10\x02\x12\x0c\n\x08\x42USINESS\x10\x03\x12\t\n\x05\x46IRST\x10\x04*E\n\x04Trip\x12\x10\n\x0cUNKNOWN_TRIP\x10\x00\x12\x0e\n\nROUND_TRIP\x10\x01\x12\x0b\n\x07ONE_WAY\x10\x02\x12\x0e\n\nMULTI_CITY\x10\x03*_\n\tPassenger\x12\x15\n\x11UNKNOWN_PASSENGER\x10\x00\x12\t\n\x05\x41\x44ULT\x10\x01\x12\t\n\x05\x43HILD\x10\x02\x12\x12\n\x0eINFANT_IN_SEAT\x10\x03\x12\x11\n\rINFANT_ON_LAP\x10\x04\x62\x06proto3') _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, globals()) -_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, "flights_pb2", globals()) +_builder.BuildTopDescriptorsAndMessages(DESCRIPTOR, 'flights_pb2', globals()) if _descriptor._USE_C_DESCRIPTORS == False: - DESCRIPTOR._options = None - _SEAT._serialized_start = 242 - _SEAT._serialized_end = 325 - _TRIP._serialized_start = 327 - _TRIP._serialized_end = 396 - _PASSENGER._serialized_start = 398 - _PASSENGER._serialized_end = 493 - _AIRPORT._serialized_start = 17 - _AIRPORT._serialized_end = 43 - _FLIGHTDATA._serialized_start = 45 - _FLIGHTDATA._serialized_end = 131 - _INFO._serialized_start = 133 - _INFO._serialized_end = 240 + + DESCRIPTOR._options = None + _SEAT._serialized_start=280 + _SEAT._serialized_end=363 + _TRIP._serialized_start=365 + _TRIP._serialized_end=434 + _PASSENGER._serialized_start=436 + _PASSENGER._serialized_end=531 + _AIRPORT._serialized_start=17 + _AIRPORT._serialized_end=43 + _FLIGHTDATA._serialized_start=45 + _FLIGHTDATA._serialized_end=169 + _INFO._serialized_start=171 + _INFO._serialized_end=278 # @@protoc_insertion_point(module_scope)