141 lines
3.2 KiB
Python
Executable file
141 lines
3.2 KiB
Python
Executable file
#!/usr/bin/env python3
|
|
|
|
from dataclasses import dataclass
|
|
import re
|
|
import json
|
|
import socketio
|
|
|
|
WS_URL = "wss://trainmap.belgiantrain.be/socket.io/"
|
|
|
|
|
|
@dataclass
|
|
class JsonLonLatCoord:
|
|
lon: float
|
|
lat: float
|
|
def __repr__(self):
|
|
return str(self.__to_json__())
|
|
def __to_json__(self):
|
|
# return f"[{self.lon:.6f}, {self.lat:.6f}]"
|
|
return [round(self.lon, 6), round(self.lat, 6)]
|
|
|
|
|
|
@dataclass
|
|
class JsonRoundingFloat:
|
|
value: float
|
|
decimals: int
|
|
def __repr__(self):
|
|
return str(self.__to_json__())
|
|
def __to_json__(self):
|
|
# return f"{{0:.{self.decimals}f}}".format(self.value)
|
|
return round(self.value, self.decimals)
|
|
|
|
|
|
class RETAIN_KEY: pass
|
|
class RETAIN_VALUE: pass
|
|
class DISCARD: pass
|
|
def transform_dict(specification: dict, d: dict):
|
|
"""
|
|
specification is a dict where the key is used to determine which action to take with the same
|
|
key in the data. The values to use in specification are:
|
|
- keys that should be discarded: DISCARD
|
|
- keys whose key and value should be retained verbatim: not present
|
|
- keys whose key and/or value should be transformed, with None denoting no change: tuple of
|
|
(str, function), or any of those two RETAIN_KEY or RETAIN_VALUE respectively
|
|
|
|
{
|
|
"""
|
|
return {
|
|
(
|
|
specification[k][0]
|
|
if specification.get(k) not in (None, DISCARD) and specification[k][0] is not RETAIN_KEY
|
|
else k
|
|
): (
|
|
specification[k][1](v)
|
|
if specification.get(k) not in (None, DISCARD) and specification[k][1] is not RETAIN_VALUE
|
|
else v
|
|
)
|
|
for k, v in d.items()
|
|
if specification.get(k) is not DISCARD
|
|
}
|
|
|
|
|
|
def camel_to_snake(s: str):
|
|
if s.islower():
|
|
return s
|
|
return re.sub(
|
|
r"(?<!^)([A-Z])", # Match an uppercase letter that is not at the start
|
|
r"_\1", # and insert an underscore
|
|
s
|
|
).lower()
|
|
|
|
|
|
def camel_to_snake_keys(data: dict):
|
|
return {
|
|
camel_to_snake(k): v
|
|
for k, v in data.items()
|
|
}
|
|
|
|
|
|
def filtered_data(data):
|
|
return {
|
|
"last_update_success": data["last_updated"]["success"],
|
|
"timestamp": data["timestamp"],
|
|
"trips": {
|
|
trip["trip_short_name"]: transform_dict({
|
|
"trip_id": DISCARD,
|
|
"trip_short_name": DISCARD,
|
|
"position": (RETAIN_KEY, lambda l: (
|
|
JsonLonLatCoord(*l)
|
|
if isinstance(l, list) else l
|
|
)),
|
|
"current_traveled_distance": (RETAIN_KEY, lambda x: (
|
|
JsonRoundingFloat(x, 1)
|
|
if isinstance(x, float) else x
|
|
)),
|
|
"fromstationdistance": ("from_station_distance", RETAIN_VALUE),
|
|
"nextstationdistance": ("next_station_distance", RETAIN_VALUE),
|
|
"trip_headsign": (RETAIN_KEY, lambda x: transform_dict({
|
|
"default": DISCARD,
|
|
"en": DISCARD,
|
|
"de": DISCARD,
|
|
}, x)),
|
|
}, camel_to_snake_keys(trip)) if isinstance(trip, dict) else trip
|
|
for trip in data["trips"].values()
|
|
},
|
|
}
|
|
|
|
|
|
class JsonEncoder(json.JSONEncoder):
|
|
def default(self, o):
|
|
if hasattr(o, "__to_json__"):
|
|
return o.__to_json__()
|
|
else:
|
|
return super().default(o)
|
|
|
|
|
|
def json_encode(data, *args, **kwargs):
|
|
return json.dumps(data, *args, cls=JsonEncoder, **kwargs)
|
|
|
|
|
|
def emit(data):
|
|
print(json_encode(filtered_data(data), ensure_ascii=False))
|
|
|
|
|
|
def main():
|
|
sio = socketio.Client()
|
|
|
|
@sio.on("event")
|
|
def handle(data):
|
|
emit(data)
|
|
|
|
sio.disconnect()
|
|
|
|
|
|
|
|
sio.connect(WS_URL, transports=["websocket"])
|
|
|
|
sio.wait()
|
|
|
|
if __name__ == "__main__":
|
|
main()
|