max_stops query parameter addition

.gitignore

@@ -1,3 +1,4 @@
 **/__pycache__/
 /main.py
 .ruff/
+.idea

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

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

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

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
     )

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 {

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,19 +17,22 @@ 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,
         *,
         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})"
+