zele-utils/functions/metro.py

from .printers import error_printer, success_printer
import geopandas, typer

def metro_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