File size: 3,232 Bytes
9be4956 6159f52 9be4956 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 |
import pandas as pd
from pandas import DataFrame
from typing import Optional
from annotation.src.utils import extract_before_parenthesis
class Flights:
def __init__(self, path="/home/user/app/database/flights/clean_Flights_2022.csv"):
self.path = path
self.data = None
self.data = pd.read_csv(self.path).dropna()[['Flight Number', 'Price', 'DepTime', 'ArrTime', 'ActualElapsedTime','FlightDate','OriginCityName','DestCityName','Distance']]
print("Flights API loaded.")
def load_db(self):
self.data = pd.read_csv(self.path).dropna().rename(columns={'Unnamed: 0': 'Flight Number'})
def run(self,
origin: str,
destination: str,
departure_date: str,
) -> DataFrame:
"""Search for flights by origin, destination, and departure date."""
results = self.data[self.data["OriginCityName"] == origin]
results = results[results["DestCityName"] == destination]
results = results[results["FlightDate"] == departure_date]
# if order == "ascPrice":
# results = results.sort_values(by=["Price"], ascending=True)
# elif order == "descPrice":
# results = results.sort_values(by=["Price"], ascending=False)
# elif order == "ascDepTime":
# results = results.sort_values(by=["DepTime"], ascending=True)
# elif order == "descDepTime":
# results = results.sort_values(by=["DepTime"], ascending=False)
# elif order == "ascArrTime":
# results = results.sort_values(by=["ArrTime"], ascending=True)
# elif order == "descArrTime":
# results = results.sort_values(by=["ArrTime"], ascending=False)
if len(results) == 0:
return "There is no flight from {} to {} on {}.".format(origin, destination, departure_date)
return results
def run_for_annotation(self,
origin: str,
destination: str,
departure_date: str,
) -> DataFrame:
"""Search for flights by origin, destination, and departure date."""
results = self.data[self.data["OriginCityName"] == extract_before_parenthesis(origin)]
results = results[results["DestCityName"] == extract_before_parenthesis(destination)]
results = results[results["FlightDate"] == departure_date]
# if order == "ascPrice":
# results = results.sort_values(by=["Price"], ascending=True)
# elif order == "descPrice":
# results = results.sort_values(by=["Price"], ascending=False)
# elif order == "ascDepTime":
# results = results.sort_values(by=["DepTime"], ascending=True)
# elif order == "descDepTime":
# results = results.sort_values(by=["DepTime"], ascending=False)
# elif order == "ascArrTime":
# results = results.sort_values(by=["ArrTime"], ascending=True)
# elif order == "descArrTime":
# results = results.sort_values(by=["ArrTime"], ascending=False)
return results
def get_city_set(self):
city_set = set()
for unit in self.data['data']:
city_set.add(unit[5])
city_set.add(unit[6]) |