code/modules/rome_gtfs_rt.py

import pandas as pd
import requests
from google.transit import gtfs_realtime_pb2
from modules.colors import IN_TRANSIT_CL, LATE_CL, ON_TIME_CL, STOPPED_CL
from modules.constants import CORS_GTFS_TRIP_UPDATES, CORS_GTFS_VEHICLE_POS
from modules.time_utils import timestamp_to_hms
from pyproj import Transformer
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry

# Vehicle Dataframe columns
VEHICLE_DF_COLUMNS = [
    "x",
    "y",
    "vehicleID",
    "tripID",
    "startTime",
    "lastUpdate",
    "currentStatus",
    "currentStatusClass",
    "statusColor",
]

# Delay Dataframe columns
DELAY_DF_COLUMNS = [
    "tripID",
    "delay",
    "delayClass",
    "delayColor",
]

VEHICLE_DF_SCHEMA = pd.DataFrame([], columns=VEHICLE_DF_COLUMNS)

DELAY_DF_SCHEMA = pd.DataFrame([], columns=DELAY_DF_COLUMNS)

FULL_DF_SCHEMA = VEHICLE_DF_SCHEMA.merge(DELAY_DF_SCHEMA, on="tripID")

# A transformer that converts coordinates from EPSG:4326 to EPSG:3857
transformer = Transformer.from_crs(4326, 3857, always_xy=True)

def get_with_retry(url, retries=5, backoff_factor=1):
    session = requests.Session()
    retry_strategy = Retry(
        total=retries,
        backoff_factor=backoff_factor,
        status_forcelist=[429, 500, 502, 503, 504],
    )
    adapter = HTTPAdapter(max_retries=retry_strategy)
    session.mount("http://", adapter)
    session.mount("https://", adapter)
    return session.get(url)

def build_url(cache_bust):
    """
    Get the current local time and build the request url
    """

    vehicle_url = CORS_GTFS_VEHICLE_POS + f"?cacheBust={cache_bust}"
    trip_url = CORS_GTFS_TRIP_UPDATES + f"?cacheBust={cache_bust}"

    return (vehicle_url, trip_url)


def get_vehicle_position(entity):
    """
    Returns the xy position of the processed entity.
    """

    coords = transformer.transform(
        entity.vehicle.position.longitude, entity.vehicle.position.latitude
    )
    return coords


def get_current_status_color(current_status):
    """
    Returns the color of the entity according to the status
    of the vehicle (In transit/Stopped).
    """

    return STOPPED_CL if current_status == 1 else IN_TRANSIT_CL


def get_current_status_class(current_status):
    """
    Returns the Vehicle current status (In transit/Stopped).
    """

    return "Stopped" if current_status == 1 else "In Transit"


def get_delay_color(delay):
    """
    Returns the color of the entity according to the delay class.
    """

    if delay <= 0:
        return ON_TIME_CL
    else:
        return LATE_CL


def get_delay_class(delay):
    """
    Returns the delay class (Late or On time).
    """

    if delay <= 0:
        return "On time"
    else:
        return "Late"


def get_vehicle_data(url):
    """Reads the vehicle position feed and returns a pandas DataFrame"""

    vehicle_feed = gtfs_realtime_pb2.FeedMessage()

    # TODO: Retry at least 5 times if the response is empty
    response = requests.get(url).content
    vehicle_feed.ParseFromString(response)

    # Entities
    vehicle_entities = vehicle_feed.entity

    positions = []
    for entity in vehicle_entities:
        # Vehicle attributes
        x, y = get_vehicle_position(entity)
        vehicle_id = entity.vehicle.vehicle.id
        trip_id = entity.vehicle.trip.trip_id.strip()
        start_time = entity.vehicle.trip.start_time
        last_update = timestamp_to_hms(entity.vehicle.timestamp)
        current_status = entity.vehicle.current_status
        current_status_class = get_current_status_class(current_status)
        vehicle_color = get_current_status_color(current_status)

        positions.append(
            [
                x,
                y,
                vehicle_id,
                trip_id,
                start_time,
                last_update,
                current_status,
                current_status_class,
                vehicle_color,
            ]
        )

    data = pd.DataFrame(positions, columns=VEHICLE_DF_COLUMNS)
    return data


def get_delay_data(url):
    """Reads the trip updates feed and returns a pandas DataFrame"""
    trip_update_feed = gtfs_realtime_pb2.FeedMessage()
#    response = requests.get(url).content
    response = get_with_retry(url).content
    trip_update_feed.ParseFromString(response)

    trip_update_entities = trip_update_feed.entity
    delays = []

    for entity in trip_update_entities:
        trip_id = entity.trip_update.trip.trip_id.strip()

        # 1. check: empty list?
        if not entity.trip_update.stop_time_update:
            print(f"Warning: No stop_time_update for trip {trip_id}. Skipping.")
            continue

        # 2. Check: does exist at least one element in list? (probably redundant)
        stop_update = entity.trip_update.stop_time_update[0]

        # ✅ 3. Check: the field 'arrival' exist?
        if not stop_update.HasField('arrival'):
            print(f"Warning: No 'arrival' field for trip {trip_id}, stop_update[0]. Skipping.")
            continue

        current_stop_arrival = stop_update.arrival
        current_stop_delay = current_stop_arrival.delay / 60  # in minuti

        delay_class = get_delay_class(current_stop_delay)
        delay_color = get_delay_color(current_stop_delay)

        delays.append([
            trip_id,
            current_stop_delay,
            delay_class,
            delay_color
        ])

    data = pd.DataFrame(delays, columns=DELAY_DF_COLUMNS)
    return data

def get_data(cache_bust):
    """
    This function reads the Roma mobilità GTFS-RT feed
    and returns a pandas DataFrame.
    """

    vehicle_url, trip_url = build_url(cache_bust)
    vehicle_data = get_vehicle_data(vehicle_url)
    delay_data = get_delay_data(trip_url)

    # Merge vehicle and delay dataframe
    full_data = vehicle_data.merge(delay_data, on="tripID")
    return full_data