zele-utils/functions/bus.py

from .printers import error_printer, success_printer
import geopandas, typer

def bus_processing(city, files):
    df_stm_arrets_sig = None
    df_stm_lignes_sig = None

    if city == "mtl":
        required_files = ["stm_arrets_sig.shp", "stm_lignes_sig.shp"]
        common_files = [file for file in files if file.name in required_files]
        if len(common_files) == 0: 
            error_printer("Incorrect file input")
            raise typer.Exit()
        
        for file in common_files:
            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_stm_arrets_sig = df_arrets_filtered
                except Exception as e:
                    error_printer(f"Failed to process stm_arrets_sig.shp: {e}")
            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_stm_lignes_sig = df_lignes_filtered
                except Exception as e:
                    error_printer(f"Failed to process stm_lignes_sig.shp: {e}")

    return df_stm_arrets_sig, df_stm_lignes_sig