1 files changed, 126 insertions, 0 deletions
diff --git a/ptv.py b/ptv.py
new file mode 100755
index 0000000..9d8b541
--- /dev/null
+++ b/ptv.py
@@ -0,0 +1,126 @@
+#!/usr/bin/python
+
+# Copyright (C) 2022 Yuchen Pei.
+#
+# ptv.py is free software: you can redistribute it and/or modify it
+# under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+
+# ptv.py is distributed in the hope that it will be useful, but
+# WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
+# Affero General Public License for more details.
+
+# You should have received a copy of the GNU Affero General Public
+# License along with ptv.py.  If not, see
+# <https://www.gnu.org/licenses/>.
+
+# PTV (Public Transport Victoria) timetable cli query tool
+
+from hashlib import sha1
+from urllib import request, error, parse
+from datetime import datetime, timezone
+from zoneinfo import ZoneInfo
+import hmac
+import json
+import pprint
+import itertools
+import os
+
+_ROUTE_TYPE = ["Train", "Tram", "Bus", "Vline", "Night Bus"]
+
+def get_url(request):
+    key = str.encode(os.environ.get('PTVKEY'))
+    dev_id = os.environ.get('PTVID')
+    request = parse.quote(request) + ('&' if ('?' in request) else '?')
+    raw = request + f'devid={dev_id}'
+    hashed = hmac.new(key, raw.encode('utf-8'), sha1)
+    signature = hashed.hexdigest().upper()
+    return f'http://timetableapi.ptv.vic.gov.au{raw}&signature={signature}'
+
+def get_json(url):
+    try:
+        document = request.urlopen(url).read().decode()
+    except error.HTTPError:
+        print("http error!")
+        return
+    return json.loads(document)
+
+def get_stops(search_result):
+    return [{'routes': stop['routes'],
+             'stop_name': stop['stop_name'],
+             'stop_id': stop['stop_id']} for stop in search_result(stops)]
+
+def get_stop_and_routes(search_result):
+    return itertools.chain.from_iterable(
+        [[(stop, route) for route in stop['routes']]
+            for stop in search_result['stops']])
+
+def search(keyword):
+    return get_json(get_url(f'/v3/search/{keyword}'))
+
+def format_stop_and_route_name(stop_and_route):
+    stop, route = stop_and_route
+    return 'Stop: {}, Route: {} {} {}'.format(
+        stop['stop_name'], _ROUTE_TYPE[route['route_type']],
+        route['route_number'], route['route_name'])
+
+def format_stop_and_route_names(stop_and_routes):
+    return '\n'.join(f'[{i}] {result}' for i, result in enumerate(
+        map(format_stop_and_route_name, stop_and_routes)))
+    return '\n'.join(f'[{i}] {result}' for i, result in enumerate(
+        [format_stop_and_route_name(
+            stop['stop_name'], _ROUTE_TYPE[route['route_type']],
+            route['route_number'], route['route_name'])
+         for stop, route in stop_and_routes]))
+
+def get_departures(stop, route):
+    return get_json(get_url(
+        '/v3/departures/route_type/{}/stop/{}/route/{}'.format(
+            route['route_type'], stop['stop_id'],
+            route['route_id'])))
+
+def parse_time(maybe_time):
+    if maybe_time:
+        return datetime.fromisoformat(maybe_time[:-1] + '+00:00')
+
+def format_time(maybe_time):
+    if maybe_time:
+        return str(maybe_time.astimezone(
+            ZoneInfo('Australia/Melbourne')))[:-6]
+
+def filter_departures(departures):
+    return [dep for dep in departures
+            if dep['scheduled_departure_utc'] and
+            parse_time(dep['scheduled_departure_utc']) >
+            datetime.now().astimezone(timezone.utc)]
+
+def get_directions(route_id):
+    return get_json(get_url(f'/v3/directions/route/{route_id}'))
+
+def get_direction_names(route_id):
+    return {dir['direction_id']: dir['direction_name'] for dir in
+            get_directions(route_id)['directions']}
+
+def format_departures(departures, direction_names):
+    return '\n'.join(['estimated: {}; scheduled: {}; direction: {}'.format(
+        format_time(parse_time(dep['estimated_departure_utc'])),
+        format_time(parse_time(dep['scheduled_departure_utc'])),
+        direction_names[dep['direction_id']])
+        for dep in filter_departures(departures['departures'])])
+
+def main():
+    stop_and_routes = list(get_stop_and_routes(search(input('Query: '))))
+    if not stop_and_routes:
+        print("No results")
+        return
+    print(format_stop_and_route_names(stop_and_routes))
+    idx = int(input('Choose a number: '))
+    print(format_stop_and_route_name(stop_and_routes[idx]))
+    print(format_departures(
+        get_departures(*stop_and_routes[idx]),
+        get_direction_names(stop_and_routes[idx][1]['route_id'])))
+
+if __name__ == "__main__":
+    main()