go-gtfs/pkg/reader/csvreader.go

package reader

import (
	"encoding/csv"
	"io"
	"strconv"

	"git.tornberg.me/go-gtfs/pkg/types"
)

func ParseAgencies(r io.Reader, callback func(types.Agency)) error {
	reader := csv.NewReader(r)
	// skip header
	if _, err := reader.Read(); err != nil {
		return err
	}
	for {
		record, err := reader.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}
		agency := types.Agency{
			AgencyID:       record[0],
			AgencyName:     record[1],
			AgencyURL:      record[2],
			AgencyTimezone: record[3],
			AgencyLang:     record[4],
		}
		callback(agency)
	}
	return nil
}

func ParseCalendars(r io.Reader, callback func(types.Calendar)) error {
	reader := csv.NewReader(r)
	if _, err := reader.Read(); err != nil {
		return err
	}
	for {
		record, err := reader.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}
		monday, _ := strconv.Atoi(record[1])
		tuesday, _ := strconv.Atoi(record[2])
		wednesday, _ := strconv.Atoi(record[3])
		thursday, _ := strconv.Atoi(record[4])
		friday, _ := strconv.Atoi(record[5])
		saturday, _ := strconv.Atoi(record[6])
		sunday, _ := strconv.Atoi(record[7])
		calendar := types.Calendar{
			ServiceID: record[0],
			Monday:    monday,
			Tuesday:   tuesday,
			Wednesday: wednesday,
			Thursday:  thursday,
			Friday:    friday,
			Saturday:  saturday,
			Sunday:    sunday,
			StartDate: record[8],
			EndDate:   record[9],
		}
		callback(calendar)
	}
	return nil
}

func ParseCalendarDates(r io.Reader, callback func(types.CalendarDate)) error {
	reader := csv.NewReader(r)
	if _, err := reader.Read(); err != nil {
		return err
	}
	for {
		record, err := reader.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}
		exceptionType, _ := strconv.Atoi(record[2])
		calendarDate := types.CalendarDate{
			ServiceID:     record[0],
			Date:          record[1],
			ExceptionType: exceptionType,
		}
		callback(calendarDate)
	}
	return nil
}

func ParseFeedInfos(r io.Reader, callback func(id string, data types.FeedInfo)) error {
	reader := csv.NewReader(r)
	if _, err := reader.Read(); err != nil {
		return err
	}
	for {
		record, err := reader.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}
		feedInfo := types.FeedInfo{
			//FeedID:            record[0],
			PublisherName: record[1],
			PublisherURL:  record[2],
			Language:      record[3],
			Version:       record[4],
		}
		callback(record[0], feedInfo)
	}
	return nil
}

func ParseRoutes(r io.Reader, callback func(types.Route)) error {
	reader := csv.NewReader(r)
	if _, err := reader.Read(); err != nil {
		return err
	}
	for {
		record, err := reader.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}
		routeType, _ := strconv.Atoi(record[4])
		route := types.Route{
			Trips:          make([]*types.Trip, 0),
			RouteID:        record[0],
			AgencyID:       record[1],
			RouteShortName: record[2],
			RouteLongName:  record[3],
			RouteType:      routeType,
			RouteURL:       record[5],
		}
		callback(route)
	}
	return nil
}

func ParseStopTimes(r io.Reader, callback func(data types.StopTime)) error {
	reader := csv.NewReader(r)
	if _, err := reader.Read(); err != nil {
		return err
	}
	for {
		record, err := reader.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}
		stopSequence, _ := strconv.Atoi(record[4])
		pickupType, _ := strconv.Atoi(record[5])
		dropOffType, _ := strconv.Atoi(record[6])
		stopTime := types.StopTime{
			TripId:        record[0],
			ArrivalTime:   types.ParseTimeString(record[1]),
			DepartureTime: types.ParseTimeString(record[2]),
			StopId:        record[3],
			StopSequence:  stopSequence,
			PickupType:    pickupType,
			DropOffType:   dropOffType,
		}
		callback(stopTime)
	}
	return nil
}

func ParseStops(r io.Reader, callback func(types.Stop)) error {
	reader := csv.NewReader(r)
	if _, err := reader.Read(); err != nil {
		return err
	}
	for {
		record, err := reader.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}
		stopLat, _ := strconv.ParseFloat(record[2], 64)
		stopLon, _ := strconv.ParseFloat(record[3], 64)
		locationType, _ := strconv.Atoi(record[4])
		stop := types.Stop{
			Trips:        make(map[string]*types.Trip),
			Transfers:    make([]*types.Transfer, 0),
			StopId:       record[0],
			StopName:     record[1],
			StopLat:      stopLat,
			StopLon:      stopLon,
			LocationType: locationType,
		}
		callback(stop)
	}
	return nil
}

func ParseTransfers(r io.Reader, callback func(types.Transfer)) error {
	reader := csv.NewReader(r)
	if _, err := reader.Read(); err != nil {
		return err
	}
	for {
		record, err := reader.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}
		transferType, _ := strconv.Atoi(record[2])
		minTransferTime, _ := strconv.Atoi(record[3])
		transfer := types.Transfer{
			FromStopId:      record[0],
			ToStopId:        record[1],
			TransferType:    transferType,
			MinTransferTime: minTransferTime,
			FromTripId:      record[4],
			ToTripId:        record[5],
		}
		callback(transfer)
	}
	return nil
}

func ParseTrips(r io.Reader, callback func(types.Trip)) error {
	reader := csv.NewReader(r)
	if _, err := reader.Read(); err != nil {
		return err
	}
	for {
		record, err := reader.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}
		trip := types.Trip{
			RouteId:       record[0],
			ServiceId:     record[1],
			TripId:        record[2],
			TripHeadsign:  record[3],
			TripShortName: record[4],
		}
		callback(trip)
	}
	return nil
}