gistfile1.py

import argparse
import re
import unittest
import datetime

from collections import defaultdict
from itertools import cycle
from functools import partial

from pymongo import MongoClient

TIMES = re.compile('(\d+:\d+)')

def main(file_name, direction_name, tt_type):
    timetable = defaultdict(list)
    with open(file_name) as f:
        for line in f.readlines():
            times = TIMES.findall(line)

            first_time = times[0]
            station_name = line[:line.index(first_time)-1]
            timetable[station_name].extend(times)

    client = MongoClient('mydocker', 49154)
    collection = client['romalido']['timetables']
    for station_name, times in timetable.iteritems():
        collection.insert({
            'station': station_name,
            'direction': direction_name,
            'type': tt_type,
            'data': times
        })

def make_times_from(timetable, current_time):
    len_timetable = len(timetable)
    loop = cycle(timetable)

    dt_maker = partial(
        datetime.datetime,
        current_time.year,
        current_time.month,
        current_time.day,
    )

    days_more = -1
    already_yielded = False
    for idx, time in enumerate(loop):
        hour, minute = map(int, time.split(':'))
        if idx % len_timetable == 0:
            days_more += 1

        if already_yielded or (hour >= current_time.hour and minute >= current_time.minute):
            new_dt = dt_maker(hour, minute) + datetime.timedelta(days=days_more)
            yield new_dt
            already_yielded = True

class TimetableTestCase(unittest.TestCase):

    def setUp(self):
        client = MongoClient('mydocker', 49154)
        self.collection = client['romalido']['timetables']

    def test_it_can_iterate(self):
        timetable = self.collection.find_one({'station': 'ACILIA', 'direction': 'towards-rome',
                                             'type': 'mon-fri'})['data']

        time = datetime.datetime(2015, 1, 18, 14, 47)
        gen = make_times_from(timetable, time)
        self.assertEquals(gen.next(), datetime.datetime(2015, 1, 18, 14, 48))
        self.assertEquals(gen.next(), datetime.datetime(2015, 1, 18, 14, 58))

    def test_it_works_for_new_day(self):
        timetable = self.collection.find_one({'station': 'ACILIA', 'direction': 'towards-rome',
                                             'type': 'mon-fri'})['data']

        time = datetime.datetime(2015, 1, 18, 23, 32)
        gen = make_times_from(timetable, time)
        self.assertEquals(gen.next(), datetime.datetime(2015, 1, 18, 23, 33))
        self.assertEquals(gen.next(), datetime.datetime(2015, 1, 18, 23, 48))
        self.assertEquals(gen.next(), datetime.datetime(2015, 1, 19, 05, 33))
        self.assertEquals(gen.next(), datetime.datetime(2015, 1, 19, 05, 48))


if __name__ == '__main__':
    #parser = argparse.ArgumentParser()
    #parser.add_argument('--direction-name', type=str, required=True)
    #parser.add_argument('--file-name', type=str, required=True)
    #parser.add_argument('--type', type=str, required=True)
    #args = parser.parse_args()
    #main(args.file_name, args.direction_name, args.type)
    unittest.main()