next_path.py

#!/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()