diff --git a/mta_manager/mta.py b/mta_manager/mta.py index 82104ee..3b114b6 100644 --- a/mta_manager/mta.py +++ b/mta_manager/mta.py @@ -4,8 +4,8 @@ import json from google.transit import gtfs_realtime_pb2 from protobuf_to_dict import protobuf_to_dict -from .train import get_train_from_dict from time import time +from train import Train class MTA(object): @@ -52,7 +52,7 @@ class MTA(object): feed = gtfs_realtime_pb2.FeedMessage() feed.ParseFromString(r.content) subway_feed = protobuf_to_dict(feed)['entity'] - trains.extend([train for train in [get_train_from_dict(train_dict) for train_dict in subway_feed] if train is not None]) + trains.extend([train for train in [Train.get_train_from_dict(train_dict) for train_dict in subway_feed] if train is not None]) return trains diff --git a/mta_manager/stop.py b/mta_manager/stop.py index f8e6ac0..2d6318d 100644 --- a/mta_manager/stop.py +++ b/mta_manager/stop.py @@ -2,10 +2,7 @@ from datetime import datetime from math import trunc -def get_stop_from_dict(obj): - if "arrival" in obj and "departure" in obj and "stop_id" in obj: - return Stop( obj["stop_id"], obj["arrival"]["time"], obj["departure"]["time"]) - return None + class Stop(object): @@ -22,3 +19,9 @@ class Stop(object): time = datetime.fromtimestamp(self.arrival_time) time_minutes = trunc(((time - now).total_seconds()) / 60) return f"stop_id:{self.id}| arr:{time_minutes}| dep:{self.departure_time}" + + @staticmethod + def get_stop_from_dict(obj): + if "arrival" in obj and "departure" in obj and "stop_id" in obj: + return Stop(obj["stop_id"], obj["arrival"]["time"], obj["departure"]["time"]) + return None diff --git a/mta_manager/train.py b/mta_manager/train.py index 9693f42..c4d66a4 100644 --- a/mta_manager/train.py +++ b/mta_manager/train.py @@ -1,17 +1,4 @@ -from .stop import get_stop_from_dict - - -def get_train_from_dict(obj): - if "trip_update" in obj and "stop_time_update" in obj["trip_update"]: - # data we need is here create object - id = obj["id"] - route = obj["trip_update"]["trip"]["route_id"] - all_stops = [get_stop_from_dict(x) for x in obj["trip_update"]["stop_time_update"]] - valid_stops = [valid_stop for valid_stop in all_stops if valid_stop is not None] - return Train(id, route, valid_stops) - else: - return None - +from stop import Stop class Train(object): def __init__(self, id, route, stops): @@ -41,3 +28,15 @@ class Train(object): def __str__(self): formatted_stops = '\n'.join([str(stop) for stop in self.stops]) return f"train_id:{self.id} | line_name:{self.route}| stops:\n {formatted_stops}" + + @staticmethod + def get_train_from_dict(obj): + if "trip_update" in obj and "stop_time_update" in obj["trip_update"]: + # data we need is here create object + id = obj["id"] + route = obj["trip_update"]["trip"]["route_id"] + all_stops = [Stop.get_stop_from_dict(x) for x in obj["trip_update"]["stop_time_update"]] + valid_stops = [valid_stop for valid_stop in all_stops if valid_stop is not None] + return Train(id, route, valid_stops) + else: + return None \ No newline at end of file