1 files changed, 73 insertions, 0 deletions
diff --git a/ptv_client.py b/ptv_client.py
new file mode 100644
index 0000000..652d87b
--- /dev/null
+++ b/ptv_client.py
@@ -0,0 +1,73 @@
+# 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/>.
+
+from hashlib import sha1
+from urllib import request, error, parse
+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 get_departures(route_type, stop_id, route_id):
+    return get_json(get_url(
+        f'/v3/departures/route_type/{route_type}/stop/{stop_id}/route/{route_id}'))
+
+def get_departures_from_stop_and_route(stop, route):
+    return get_departures(route['route_type'], stop['stop_id'],
+                          route['route_id'])
+
+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 get_route_type(idx):
+    return _ROUTE_TYPE[idx]