zele-utils/functions/metro.py

34 lines
1.4 KiB
Python
Raw Normal View History

2024-09-15 17:04:21 -04:00
from .printers import error_printer, success_printer
import geopandas, typer
2024-09-07 18:13:36 -04:00
2024-09-15 17:04:21 -04:00
def metro_processing(city, files):
2024-09-07 18:13:36 -04:00
2024-09-15 17:04:21 -04:00
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.stem in required_files]
if len(common_files) == 0:
error_printer("Incorrect file input")
raise typer.Exit()
for file in common_files:
if file == 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 == 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