Skip to content

Instantly share code, notes, and snippets.

@hotpxl
Created October 20, 2019 01:32
Show Gist options
  • Save hotpxl/d2c9b5c9e502a37b3eae2518e0810587 to your computer and use it in GitHub Desktop.
Save hotpxl/d2c9b5c9e502a37b3eae2518e0810587 to your computer and use it in GitHub Desktop.
#!/usr/bin/env python
# -*- encoding: utf-8 -*-
from __future__ import print_function
import csv
import re
import datetime
def get_service_ids(today):
service_ids = set()
with open("calendar.txt") as f:
reader = csv.reader(f)
next(reader)
for row in reader:
start_date = int(row[-2])
end_date = int(row[-1])
if (
start_date
<= (10000 * today.year + 100 * today.month + today.day)
<= end_date
and row[today.weekday() + 2] == "1"
):
service_ids.add(row[0])
return service_ids
def get_stop_ids(name):
r = re.compile(name, re.IGNORECASE)
stop_ids = set()
with open("stops.txt") as f:
reader = csv.reader(f)
next(reader)
for row in reader:
if r.search(row[3]):
stop_ids.add(row[0])
return stop_ids
def get_trips(today):
service_ids = get_service_ids(today)
trips = {}
with open("trips.txt") as f:
reader = csv.reader(f)
next(reader)
for row in reader:
if row[1] in service_ids:
trips[row[2]] = row[4]
return trips
def parse_time(today, t):
h, m, s = map(int, t.split(":"))
return datetime.datetime(today.year, today.month, today.day) + datetime.timedelta(
hours=h, minutes=m, seconds=s
)
def filter_trips(today, trip_ids, from_stop_ids, to_stop_ids):
trips = {}
stop_times = []
with open("stop_times.txt") as f:
reader = csv.reader(f)
next(reader)
for row in reader:
trip_id = row[0]
departure_time = parse_time(today, row[2])
stop_id = row[3]
if trip_id not in trip_ids:
continue
if stop_id in from_stop_ids:
outgoing = True
elif stop_id in to_stop_ids:
outgoing = False
else:
continue
if trip_id not in trips:
trips[trip_id] = (stop_id, departure_time)
continue
(other_stop_id, other_departure_time) = trips.pop(trip_id)
assert other_stop_id != stop_id
if (other_departure_time < departure_time) == outgoing:
# Wrong direction
continue
stop_times.append(
(trip_id, departure_time if outgoing else other_departure_time)
)
return stop_times
def main():
now = datetime.datetime.now()
trips = []
for look_ahead in range(2):
today = now.date() + datetime.timedelta(look_ahead)
all_trips = get_trips(today)
trip_ids = set(all_trips.keys())
from_stop_ids = get_stop_ids("world trade center")
to_stop_ids = get_stop_ids("exchange place")
trips.extend(
[
(all_trips[i[0]], i[1])
for i in filter_trips(today, trip_ids, from_stop_ids, to_stop_ids)
if now <= i[1]
]
)
for line, time in sorted(trips, key=lambda i: i[1])[:3]:
print("{}\t{}".format(line, time))
if __name__ == "__main__":
main()
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment