transport-accessibility/transport_accessibility/pt_map/query.py

"""
Query
-----

Module to handle database IO while abstracting the specific SQL and django model filtering.

Functions
---------
"""
from pt_map.models import *
from django.db import models
from .class_names import *
from django.http import HttpRequest
from datetime import datetime, date

def get_field_names(model: models.Model) -> list[str]:
    """
    Given a model, returns a list of the name strings of all the model's fields.
    """
    return [field.name for field in model._meta.fields]

def get_pks_from_get(req_get: dict[str, str]) -> dict[str, list[models.Model]]:
    """
    Extract primary keys from a request.GET dict-like and find the corresponding classes.

    Parameters
    ----------
    req_get : dict[str, str]
        dict-like object, a HTTP Requests GET data.

    Returns
    -------
    dict[str, list[str]]
        dict mapping a model to a list of id fields passed with GET
    """
    result = {}
    for k in req_get.keys():
        if k in classes_by_primary_keys.keys():
            result[classes_by_primary_keys[k]] = req_get.getlist(k)
    if not result:
        raise ValueError("No pks found.")
    return result

def get_obj_by_pk(mdl: models.Model, pks: list[str]) -> list[mdl.Model]:
    """
    Given a model, and a list of corresponding primary keys, return a list of objects of the given model identified by the given primary keys.

    Parameters
    ----------
    mdl: models.Model
        Model class to look for
    pks: list[str]
        primary keys of the objects to return

    Returns
    -------
    list[mdl]
        Objects corresponding to primary keys in pk.

    Raises
    ------
    mdl.DoesNotExist
        If at least one object from the list of pks could not be found.
    """
    return [obj for obj in [mdl.objects.get(**{primary_keys[mdl]: pk}) for pk in pks] if obj]
        
def obj_from_get(req_get: dict[str,str]) -> dict[cls, list[models.Model]]:
    """
    Given the GET data of a HTTP Request, return a dict with the requested model classes as keys and lists of the requested model objects as values.
    """
    return {mdl: get_obj_by_pk(mdl, keys) for mdl, keys in get_pks_from_get(req_get).items()}

def get_timetable(r: pt_map.models.Route, trips_r: list[pt_map.models.Trip], stop_sequence: list[str]):
    """ 
    Given a pt_map.models.Route, calculate the timetable for all its stops.

    Parameters
    ----------
    r : pt_map.models.Route
        Route, the timetable should be calculated for
    trips_r : list(pt_map.Trip)
        List of trips travelling on the Route r
    stop_sequence : list(str)
        List of stop_ids the Route r serves. Currently the first trip is taken as reference for stops and sequence.

    Returns
    -------
    dict{"stop_sequence": list(str), "stop_times": dict[str, list(str)]}
        Dict containing two elements:
            "stop_sequence" : list(str)
                list of stop_ids the route serves
            "stop_times" : dict(str, list(str))
                dict mapping stop_ids from stop_sequence to time strings the route is serving the stop at
    """
    timetable = {"stop_sequence": stop_sequence}
    sts = {}
    for stop in stop_sequence:
        times = []
        for t in trips_r:
            for st in StopTime.objects.filter(trip_id=t.trip_id):
                times.append(st.departure_time.strftime("%H:%M")) 
        sts[stop] = times
    timetable["stop_times"] =  sts
    return timetable

def get_all_stops() -> dict[str, dict[str,str]]:
    """
    Return all Stop object stored in the database.

    Representation of the result:
    dict:
        {
            stop_id (str): {
                    'stop_name': pt_map.models.Stop.stop_name,
                    'stop_lat': pt_map.models.Stop.stop_lat,
                    'stop_lon': pt_map.models.Stop.stop_lon,
                }
        }
    """
    return {s.stop_id: {name: getattr(s, name) for name in ['stop_name', 'stop_lat', 'stop_lon']} for s in Stop.objects.all()} 

def get_all_routes() -> list[dict[str, str]]:
    """
    Return a list of all Route objects found in the database.

    Representation of the result:
    list:
        [
                {
                    'route_id': pt_map.models.Route.route_id, 
                    'route_type': pt_map.models.Route.route_type,
                    'route_name': pt_map.models.Route.route_short_name if set else pt_map.models.Route.route_long_name,
                    'agency_id': pt_map.models.Route.agency_id.agency_id,
                }
        ]
    """
    route_name = lambda r : r.route_short_name if r.route_short_name else r.route_long_name
    return [{"route_id": r.route_id, "route_type": r.route_type, "route_name": route_name(r), "agency_id": r.agency_id.agency_id} for r in Route.objects.all()]

def get_trips(routes: list[pt_map.models.Route]) -> dict[str, list[pt_map.models.Trip]]:
    """
    Return a list of all Trips associated with a Route in the argument.
    
    Parameters
    ----------
    routes: list[str]
        List of primary keys for the Routes to search the trips for.

    Returns
    -------
    dict[str, list[pt_map.models.Trip]]
        Keys: route_ids from parameter
        Values: lists of corresponding trip objects.
    """
    return {r["route_id"]: [t for t in Trip.objects.filter(route_id_id=r["route_id"])] for r in routes} 

def get_stop_sequences(routes: list[pt_map.models.Route], trips: dict[str,list[pt_map.models.Trip]]=None) -> dict[str, list[str]]:
    """
    For all given routes, return a list of stops in the order of appearance along the route. The first trip in the list of trips is used to define the sequence.

    Parameters
    ----------
    routes: list[pt_map.models.Route]
        List of pt_map.models.Route to find stop sequences for
    trips: dict[str,list[pt_map.models.Trip]]
        List of at least one trip for each Route in routes. If none, all are calculated and the first used for the sequence.

    Returns
    -------
    dict[str, list[str]]
        Keys: route_ids
        Values: Lists of stop_ids in the order of appearance in the first Trip in the routes' trips list given.
    """
    if not trips:
        trips = get_trips(routes)
    stop_sequences = {}
    for r in routes:
        seq = []
        t = trips[r["route_id"]]
        for s in StopTime.objects.filter(trip_id_id__exact=t[0].trip_id):
            seq.append(s)
        stop_sequences[r["route_id"]] = [s.stop_id.stop_id for s in sorted(seq, key=lambda st : st.stop_sequence)]
    return stop_sequences