|
| 1 | +# ==> Configuration. |
| 2 | + |
| 3 | +# The URL of the GTFS feed. |
| 4 | +# For the NYC MTA feeds, see: https://api.mta.info/#/subwayRealTimeFeeds |
| 5 | +GTFS_FEED_URL = "https://api-endpoint.mta.info/Dataservice/mtagtfsfeeds/nyct%2Fgtfs" |
| 6 | + |
| 7 | +# The ID of the stop to watch. |
| 8 | +# For the NYC MTA, see: http://web.mta.info/developers/data/nyct/subway/google_transit.zip |
| 9 | +STOP_ID = "123S" |
| 10 | + |
| 11 | +# The color of the bullet to use for arrivals. |
| 12 | +BULLET_COLOR = "#ee352e" |
| 13 | + |
| 14 | +# Number of hardcoded iterations to simulate an infinite loop, since Starlark |
| 15 | +# doesn't support infinite loops. Used for decoding protobuf messages and such, |
| 16 | +# which can support an arbitrary number of fields, etc. |
| 17 | +MAX_ITERATIONS = 64 |
| 18 | + |
| 19 | +# ==> Entrypoint. |
| 20 | + |
| 21 | +load("render.star", "render") |
| 22 | +load("http.star", "http") |
| 23 | +load("time.star", "time") |
| 24 | + |
| 25 | +def main(): |
| 26 | + feed = gtfs_get_feed(GTFS_FEED_URL) |
| 27 | + arrivals = gtfs_get_upcoming_arrivals(STOP_ID, feed) |
| 28 | + return render.Root(render_arrivals(arrivals)) |
| 29 | + |
| 30 | +# ==> Rendering. |
| 31 | + |
| 32 | +def render_arrivals(arrivals): |
| 33 | + arrivals = [arrivals[:4], arrivals[4:]] |
| 34 | + return render.Row( |
| 35 | + children = [ |
| 36 | + render.Padding( |
| 37 | + child = render.Column( |
| 38 | + children = [render_arrival(arrival) for arrival in arrivals], |
| 39 | + ), |
| 40 | + pad = (1, 0, 2, 0), |
| 41 | + ) |
| 42 | + for arrivals in arrivals |
| 43 | + ], |
| 44 | + expanded = True, |
| 45 | + main_align = "space_between", |
| 46 | + ) |
| 47 | + |
| 48 | +def render_arrival(arrival): |
| 49 | + return render.Row( |
| 50 | + children = [ |
| 51 | + render_bullet(arrival.route_id), |
| 52 | + render.Padding(child = render_eta(arrival.eta), pad = (2, 0, 0, 0)), |
| 53 | + ], |
| 54 | + ) |
| 55 | + |
| 56 | +def render_bullet(route): |
| 57 | + return render.Circle( |
| 58 | + color = BULLET_COLOR, |
| 59 | + diameter = 7, |
| 60 | + child = render.Text(route, font = "tom-thumb"), |
| 61 | + ) |
| 62 | + |
| 63 | +def render_eta(eta): |
| 64 | + # Truncating rounding is useful as we want to miss in the "arriving too |
| 65 | + # soon" direction. |
| 66 | + eta = int((eta - time.now()).minutes) |
| 67 | + if eta < 1: |
| 68 | + return render.Text("now", font = "tb-8", color = "#ffa500") |
| 69 | + else: |
| 70 | + return render.Text("{}m".format(eta), font = "tb-8") |
| 71 | + |
| 72 | +# ==> GTFS-specific decoding. |
| 73 | +# |
| 74 | +# We hand roll a protobuf decoder here because Starlark doesn't have built-in |
| 75 | +# support for parsing protobufs (nor a package system). The alternative would be |
| 76 | +# to run a service somewhere that converts the protobuf API to a JSON API, but |
| 77 | +# that would be annoying to maintain. While this was painful to write it is |
| 78 | +# likely to continue to operate without maintenance for the foreseeable future. |
| 79 | + |
| 80 | +# Field number constants. Extracted from https://gtfs.org/documentation/realtime/proto |
| 81 | +# on 24 November 2024. |
| 82 | +FEED_MESSAGE_ENTITY_FIELD_NUMBER = 2 |
| 83 | +FEED_ENTITY_TRIP_UPDATE_FIELD_NUMBER = 3 |
| 84 | +TRIP_UPDATE_TRIP_FIELD_NUMBER = 1 |
| 85 | +TRIP_UPDATE_STOP_TIME_UPDATE_FIELD_NUMBER = 2 |
| 86 | +TRIP_DESCRIPTOR_ROUTE_ID_FIELD_NUMBER = 5 |
| 87 | +STOP_TIME_UPDATE_STOP_ID_FIELD_NUMBER = 4 |
| 88 | +STOP_TIME_UPDATE_ARRIVAL_FIELD_NUMBER = 2 |
| 89 | +STOP_TIME_EVENT_ARRIVAL_TIME_FIELD_NUMBER = 2 |
| 90 | + |
| 91 | +def gtfs_get_feed(url): |
| 92 | + return http.get(url, ttl_seconds = 5).body() |
| 93 | + |
| 94 | +def gtfs_get_upcoming_arrivals(stop_id, reader): |
| 95 | + upcoming_arrivals = [] |
| 96 | + feed_message, _ = proto_decode_message(reader) |
| 97 | + for feed_entity in feed_message.get(FEED_MESSAGE_ENTITY_FIELD_NUMBER, []): |
| 98 | + feed_entity, _ = proto_decode_message(feed_entity) |
| 99 | + for trip_update in feed_entity.get(FEED_ENTITY_TRIP_UPDATE_FIELD_NUMBER, []): |
| 100 | + trip_update, _ = proto_decode_message(trip_update) |
| 101 | + trip_descriptors = trip_update.get(TRIP_UPDATE_TRIP_FIELD_NUMBER, [""]) |
| 102 | + trip_descriptor, _ = proto_decode_message(trip_descriptors[0]) |
| 103 | + route_ids = trip_descriptor.get(TRIP_DESCRIPTOR_ROUTE_ID_FIELD_NUMBER, [""]) |
| 104 | + route_id = route_ids[0] |
| 105 | + for stop_time_update in trip_update.get(TRIP_UPDATE_STOP_TIME_UPDATE_FIELD_NUMBER, []): |
| 106 | + stop_time_update, _ = proto_decode_message(stop_time_update) |
| 107 | + stop_ids = stop_time_update.get(STOP_TIME_UPDATE_STOP_ID_FIELD_NUMBER, []) |
| 108 | + if stop_ids != [STOP_ID]: |
| 109 | + continue |
| 110 | + arrivals = stop_time_update.get(STOP_TIME_UPDATE_ARRIVAL_FIELD_NUMBER, [""]) |
| 111 | + arrival, _ = proto_decode_message(arrivals[0]) |
| 112 | + arrival_times = arrival.get(STOP_TIME_EVENT_ARRIVAL_TIME_FIELD_NUMBER, [0]) |
| 113 | + arrival_time = time.from_timestamp(arrival_times[0]) |
| 114 | + if arrival_time > time.now(): |
| 115 | + upcoming_arrivals.append(struct(route_id = route_id, eta = arrival_time)) |
| 116 | + |
| 117 | + return sorted(upcoming_arrivals, key = lambda x: x.eta) |
| 118 | + |
| 119 | +# ==> Generic Protobuf decoding. |
| 120 | + |
| 121 | +PROTO_WIRE_TYPE_VARINT = 0 |
| 122 | +PROTO_WIRE_TYPE_I64 = 1 |
| 123 | +PROTO_WIRE_TYPE_LEN = 2 |
| 124 | +PROTO_WIRE_TYPE_SGROUP = 3 |
| 125 | +PROTO_WIRE_TYPE_EGROUP = 4 |
| 126 | +PROTO_WIRE_TYPE_I32 = 5 |
| 127 | + |
| 128 | +def proto_decode_message(reader): |
| 129 | + """Decode a single message from a protobuf reader. |
| 130 | +
|
| 131 | + Returns a dict mapping field numbers to lists of field values. A required |
| 132 | + field will always appear in the output dict and its list will have exactly |
| 133 | + one value. Optional and repeated fields may or may not appear in the output |
| 134 | + dict. When an optional field appears, it will have exactly one value in its |
| 135 | + list. When a repeated field appears, it will have one or more values in its |
| 136 | + list. |
| 137 | + """ |
| 138 | + out = dict() |
| 139 | + for _ in range(MAX_ITERATIONS): |
| 140 | + if len(reader) == 0: |
| 141 | + break |
| 142 | + field_number, field_value, reader = proto_decode_field(reader) |
| 143 | + if field_number not in out: |
| 144 | + out[field_number] = [] |
| 145 | + out[field_number].append(field_value) |
| 146 | + return out, reader |
| 147 | + |
| 148 | +def proto_decode_field(reader): |
| 149 | + varint, reader = proto_decode_varint(reader) |
| 150 | + field_number = varint >> 3 |
| 151 | + wire_type = varint & 0x07 |
| 152 | + if wire_type == PROTO_WIRE_TYPE_VARINT: |
| 153 | + field_value, reader = proto_decode_varint(reader) |
| 154 | + elif wire_type == PROTO_WIRE_TYPE_LEN: |
| 155 | + field_value, reader = proto_decode_len(reader) |
| 156 | + # WARNING: many other wire types ignored, as they do not appear in the |
| 157 | + # GTFS protobufs. |
| 158 | + |
| 159 | + else: |
| 160 | + fail("proto_decode_field: unknown wire type: {}".format(wire_type)) |
| 161 | + return field_number, field_value, reader |
| 162 | + |
| 163 | +def proto_decode_varint(reader): |
| 164 | + out = 0 |
| 165 | + shift = 0 |
| 166 | + for _ in range(MAX_ITERATIONS): |
| 167 | + byte, reader = proto_next_byte(reader) |
| 168 | + out += (byte & 0x7f) << shift |
| 169 | + shift += 7 |
| 170 | + if (byte & 0x80) == 0: |
| 171 | + break |
| 172 | + return out, reader |
| 173 | + |
| 174 | +def proto_decode_len(reader): |
| 175 | + len, reader = proto_decode_varint(reader) |
| 176 | + out = reader[:len] |
| 177 | + reader = reader[len:] |
| 178 | + return out, reader |
| 179 | + |
| 180 | +def proto_next_byte(reader): |
| 181 | + out = reader.elem_ords()[0] |
| 182 | + reader = reader[1:] |
| 183 | + return out, reader |
0 commit comments