zele-utils/functions/bus.py
2024-09-16 08:38:38 -04:00

33 lines
1.3 KiB
Python

from .printers import error_printer, success_printer
import geopandas, typer
def bus_processing(city, file):
df = None
file_name = file.stem
if city == "mtl":
required_files = ["stm_arrets_sig", "stm_lignes_sig"]
if file_name not in required_files:
error_printer("Incorrect file input")
return None
if file_name == required_files[0]:
try:
df_arrets = geopandas.read_file(file)
df_arrets_filtered = df_arrets[~df_arrets['stop_url'].str.contains('metro', case=False, na=False)]
df_arrets_filtered = df_arrets_filtered.rename(columns={'geometry': 'coordinates'})
df = df_arrets_filtered
except Exception as e:
error_printer(f"Failed to process stm_arrets_sig.shp: {e}")
raise typer.Exit()
elif file_name == required_files[1]:
try:
df_lignes = geopandas.read_file(file)
df_lignes_filtered = df_lignes[~df_lignes['route_name'].str.contains('Ligne', na=False)]
df = df_lignes_filtered
except Exception as e:
error_printer(f"Failed to process stm_lignes_sig.shp: {e}")
raise typer.Exit()
return df