File size: 2,079 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 |
import pandas as pd
from pandas import DataFrame
from typing import Optional
from annotation.src.utils import extract_before_parenthesis
class Restaurants:
def __init__(self, path="/home/user/app/database/restaurants/clean_restaurant_2022.csv"):
self.path = path
self.data = pd.read_csv(self.path).dropna()[['Name','Average Cost','Cuisines','Aggregate Rating','City']]
print("Restaurants loaded.")
def load_db(self):
self.data = pd.read_csv(self.path).dropna()
def run(self,
city: str,
) -> DataFrame:
"""Search for restaurant ."""
results = self.data[self.data["City"] == city]
# results = results[results["date"] == date]
# if price_order == "asc":
# results = results.sort_values(by=["Average Cost"], ascending=True)
# elif price_order == "desc":
# results = results.sort_values(by=["Average Cost"], ascending=False)
# if rating_order == "asc":
# results = results.sort_values(by=["Aggregate Rating"], ascending=True)
# elif rating_order == "desc":
# results = results.sort_values(by=["Aggregate Rating"], ascending=False)
if len(results) == 0:
return "There is no restaurant in this city."
return results
def run_for_annotation(self,
city: str,
) -> DataFrame:
"""Search for restaurant ."""
results = self.data[self.data["City"] == extract_before_parenthesis(city)]
# results = results[results["date"] == date]
# if price_order == "asc":
# results = results.sort_values(by=["Average Cost"], ascending=True)
# elif price_order == "desc":
# results = results.sort_values(by=["Average Cost"], ascending=False)
# if rating_order == "asc":
# results = results.sort_values(by=["Aggregate Rating"], ascending=True)
# elif rating_order == "desc":
# results = results.sort_values(by=["Aggregate Rating"], ascending=False)
return results |