libtrans/src/services/TriMet.py

"""
TriMet API Wrapper
Author(s): Arris Kathery <whotookelburg@hotmail.com>

NOTE: An API Key is always required to use TriMet's API. Learn more here: https://developer.trimet.org/why_an_appid.shtml, and acquire one here: https://developer.trimet.org/appid/registration/
"""

from .. import parse_timestamp # ? this creates an error with pylint, ignore this.

import json
import requests
from logging import info, warn


class TriMetAPI:

    def __init__(self, api_key=None, secure=True) -> None:
        if api_key:
            self.api_key = api_key
        else:
            warn("The API Key was never set!")

        if secure:
            self.secure = secure
        else:
            warn("Using HTTP frontend!")

        return None

    def verify(self) -> bool:
        if self.api_key:
            info("api_key exists!")
        else:
            warn("An API Key is required to use Trimet's online services.")
            return False

        if len(self.api_key) == 25:
            info("api_key is correct length!")
            return True
        else:
            warn("api_key is at incorrect length!")
            return False

    def set_api_key(self, api_key) -> bool:
        self.api_key = api_key
        return True


class TriMetStop:
    def __init__(self) -> None:
        return None


class TriMetVehicle:

    def __init__(self, json_data=None) -> None:
        if json_data:
            self.load_from_json(json_data)
        return None

    def load_from_id(
        self, api: TriMetAPI, vehicle_id: int, api_verify: bool = True
    ) -> (
        bool
    ):  #  this is basically a glorified wrapper for load_from_json() but very lazy. also,
        # *api_verify can skip checking the api key, incase it was checked before.
        if not api.verify() and not api_verify:
            raise RuntimeError(
                "A valid API Key is required to use Trimet's online services."
            )
        else:
            info(f"loading vehicle info from id {vehicle_id}")
            return self.load_from_json(
                requests.get(
                    f"https://developer.trimet.org/ws/v2/vehicles?appID={api.api_key}&ids={vehicle_id}"
                ).json()["resultSet"]["vehicle"][0]
            )

    def load_from_json(
        self, object
    ) -> bool:  # ! INTERNAL USAGE ONLY!!! DO NOT CALL OTHERWISE!!!
        from ansiconverter import HEXtoRGB

        self.import_data = object
        self.raw_data = json.dumps(object, indent=0)

        self.id = object["vehicleID"]

        self.route_color = HEXtoRGB(f"#{object['routeColor']}")

        self.expiration = parse_timestamp(object["expires"], 0)

        self.type = None  # TODO: Identify exact model. Model should be revealed by vehicle number.
        if object["type"] == "rail":

            if object["routeSubType"] == "Light Rail":
                self.type = ["light-rail", "TriMet MAX Light Rail Vehicle"]
            else:  # Cannot check for WES or Portland Streetcar, as they are not exposed by TriMet's API.
                self.type = ["rail", "Unknown TriMet Rail Vehicle"]

        else:  # TODO: Identify LIFT and FX2
            self.type = ["bus", "TriMet Bus Vehicle"]

        self.sign = [object["signMessage"], object["signMessageLong"]]

        return True

    def refresh(
        self, api: TriMetAPI
    ) -> (
        bool
    ):  # this is a very lazy method that is just a wrapper for load_from_id(). see line #53
        if not api.verify():
            raise RuntimeError(
                "A valid API Key is required to use Trimet's online services."
            )
        else:
            return self.load_from_id(api, self.id, False)


def get_vehicles(api: TriMetAPI) -> dict:
    if not api.api_key or not api.verify():
        raise RuntimeError(
            "A valid API Key is required to use Trimet's online services."
        )
    elif api.verify():
        resp = requests.get(
            f"https://developer.trimet.org/ws/v2/vehicles?appID={api.api_key}"
        )
        result = {}

        from datetime import datetime

        result["resultTime"] = parse_timestamp(resp.json()["resultSet"]["queryTime"], 0)
        del datetime

        result["vehicles"] = []
        for vm in resp.json()["resultSet"]["vehicle"]:
            result["vehicles"].append(TriMetVehicle(json_data=vm))

        return result
    return {}