ArangoDB: Transform GTFS to Vertices and Edges for ArangoDB Importer

edges.py

import csv


collectionName = 'vertices'
gtfsStopTimes = 'gtfs/stop_times.txt'
edgesFile = 'graph/edges.json'


def readAsJson(filePath):
    i = open(filePath, 'r')
    return csv.DictReader(i)


def writeAsJson(filePath, data):
    o = open(filePath, 'w')
    o.write(data)


def calculateDuration(departure, arrival):
    departure = int(departure.replace(":", ""))
    arrival = int(arrival.replace(":", ""))
    return (arrival - departure) / 100


def makeEdges(data):
    tmp = ''
    try:
        for item in data:
            next = data.next()
            if int(next['stop_sequence']) == int(item['stop_sequence']) + 1:
                edge = {}
                edge['_from'] = collectionName + '/' + item['stop_id']
                edge['_to'] = collectionName + '/' + next['stop_id']
                edge['duration'] = calculateDuration(
                    item['departure_time'],
                    next['arrival_time']
                )
                tmp = tmp + ('%s\n' % edge)
    except StopIteration:
        tmp = tmp.replace("'", '"')
        return tmp


data = readAsJson(gtfsStopTimes)
edges = makeEdges(data)
writeAsJson(edgesFile, edges)

vertices.py

import csv


exampleVertex = {
    "_id": "stops/201477134",   # ArangoDB specific
    "_rev": "201477134",        # ArangoDB specific
    "_key": "201477134",        # ArangoDB specific
    "stop_code": "",
    "parent_station": "",
    "stop_url": "",
    "stop_desc": "",
    "location_type": "0",
    "zone_id": "",
    "stop_lat": "52.4413430",
    "stop_lon": "13.6996390",
    "stop_id": "9183504",
    "stop_name": "Petershagener Weg (Berlin)"
}

collectionName = 'vertices'
gtfsStops = 'gtfs/stops.txt'
verticesFile = 'graph/vertices.json'


def readAsJson(filePath):
    i = open(filePath, 'r')
    return csv.DictReader(i)


def writeAsJson(filePath, data):
    o = open(filePath, 'w')
    o.write(data)


def makeVertices(data):
    tmp = ''
    for item in data:
        item['_id'] = collectionName + '/' + item['stop_id']
        item['_rev'] = item['stop_id']
        item['_key'] = item['stop_id']

        tmp = tmp + ('%s\n' % item)

    tmp = tmp.replace("'", '"')
    return tmp


data = readAsJson(gtfsStops)
vertices = makeVertices(data)
writeAsJson(verticesFile, vertices)