File size: 1,215 Bytes
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 |
import pandas as pd
from pandas import DataFrame
from typing import Optional
from annotation.src.utils import extract_before_parenthesis
class Attractions:
def __init__(self, path="../database/attractions/attractions.csv"):
self.path = path
self.data = pd.read_csv(self.path).dropna()[['Name','Latitude','Longitude','Address','Phone','Website',"City"]]
print("Attractions loaded.")
def load_db(self):
self.data = pd.read_csv(self.path)
def run(self,
city: str,
) -> DataFrame:
"""Search for Accommodations by city and date."""
results = self.data[self.data["City"] == city]
# the results should show the index
results = results.reset_index(drop=True)
if len(results) == 0:
return "There is no attraction in this city."
return results
def run_for_annotation(self,
city: str,
) -> DataFrame:
"""Search for Accommodations by city and date."""
results = self.data[self.data["City"] == extract_before_parenthesis(city)]
# the results should show the index
results = results.reset_index(drop=True)
return results |