Files
kiwi-budapest/regiojet.py
2022-03-05 15:50:24 +01:00

120 lines
4.0 KiB
Python

import requests
import json
import pprint
import argparse
from datetime import date, datetime, time, timedelta
import redis
from slugify import slugify
from typing import Optional
from redis import Redis
url = "https://brn-ybus-pubapi.sa.cz"
tarrifs_path = "/restapi/consts/tariffs"
location_path = "/restapi/consts/locations"
route_path = "/restapi/routes/search/simple"
redis_host = "redis.pythonweekend.skypicker.com"
#redis_host = "localhost"
def store_dict(redis: Redis, key: str, value: dict) -> None:
redis.set(key, json.dumps(value, default=json_serial))
def retrieve_dict(redis: Redis, key: str) -> Optional[dict]:
maybe_value = redis.get(name=key)
if maybe_value is None:
return None
return json.loads(maybe_value)
def json_serial(obj):
"""JSON serializer for objects not serializable by default json code"""
if isinstance(obj, timedelta):
return str(obj)
if isinstance(obj, (datetime, date)):
return obj.isoformat()
raise TypeError("Type %s not serializable" % type(obj))
def search_locations(country,city):
for location in locations:
if country == location['country']:
for _city in location['cities']:
if city == _city['name']:
return _city
def search_connection(from_station, to_station, tariff_type, to_location_type, from_location_type, departure):
r = requests.get(url + route_path, params={"tariffs_type": tariff_type,
"toLocationId": to_station["id"],
"fromLocationId": from_station["id"],
"fromLocationType": from_location_type,
"toLocationType": to_location_type,
"departureDate": departure})
routes = json.loads(r.content)
surname = "jaro"
source = from_station["name"]
destination = to_station["name"]
key1 = F"{surname}:journey"
key2 = slugify(source)
key3 = slugify(destination)
key4 = departure
key = ':'.join((key1, key2 + '_' + key3 + '_' + key4))
journey = retrieve_dict(redisdb, key)
if journey is not None:
return journey
routes_ret = []
for route in routes['routes']:
ret = {}
ret["departure_datetime"] = datetime.fromisoformat(route["departureTime"])
ret["arrival_datetime"] = datetime.fromisoformat(route["arrivalTime"])
ret["source"] = from_station["name"]
ret["destination"] = to_station["name"]
ret["source_id"] = from_station["id"]
ret["destination_id"] = to_station["id"]
ret["free_seats"] = route["freeSeatsCount"]
ret["carrier"] = "REGIOJET"
ret["type"] = route["vehicleTypes"][0]
ret["fare"] = {"amount": route["priceFrom"], "currency": "EUR"}
routes_ret.append(ret)
store_dict(redisdb, key, routes_ret)
return routes_ret
parser = argparse.ArgumentParser(
description='Search some connection') # use of ArgumentParser against of simple OptionParser
parser.add_argument("origin")
parser.add_argument("destination")
parser.add_argument("departure")
args = parser.parse_args()
redisdb = Redis(host=redis_host, port=6379, db=0, decode_responses=True)
tariffs = retrieve_dict(redisdb,'jaro:REGIOJET:tariffs')
if tariffs is None:
r = requests.get(url + tarrifs_path)
tariffs = json.loads(r.content)
store_dict(redisdb, 'jaro:REGIOJET:tariffs', tariffs)
locations = retrieve_dict(redisdb,'jaro:REGIOJET:locations')
if locations is None:
r = requests.get(url + location_path)
locations = json.loads(r.content)
store_dict(redisdb, 'jaro:REGIOJET:locations', locations)
city_from = search_locations('Czech Republic', args.origin)
city_to = search_locations('Czech Republic', args.destination)
#pprint.pp(city_from)
#pprint.pp(city_to)
ret = search_connection(city_from, city_to, 'REGULAR', 'CITY', 'CITY', args.departure)
print(json.dumps(ret, indent=4, default=json_serial, sort_keys=False))