Skip to content

Commit 72c228b

Browse files
jamespfennelljamespfennell
committed
Add journal command
Also copies over more journaling code from the subwaydata.nyc repo, specifically the code for exporting to CSV.
1 parent f3478c1 commit 72c228b

File tree

8 files changed

+8738
-35
lines changed

8 files changed

+8738
-35
lines changed

‎cmd/gtfs.go‎

Lines changed: 104 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -3,20 +3,33 @@ package main
33
import (
44
"fmt"
55
"os"
6+
"path/filepath"
67
"strings"
78
"time"
89

910
"github.com/fatih/color"
1011
"github.com/jamespfennell/gtfs"
1112
"github.com/jamespfennell/gtfs/extensions/nyctalerts"
1213
"github.com/jamespfennell/gtfs/extensions/nycttrips"
14+
"github.com/jamespfennell/gtfs/journal"
1315
"github.com/urfave/cli/v2"
1416
)
1517

1618
func main() {
1719
app := &cli.App{
1820
Name: "GTFS parser",
1921
Usage: "parse GTFS static and realtime feeds",
22+
Flags: []cli.Flag{
23+
&cli.BoolFlag{
24+
Name: "verbose",
25+
Aliases: []string{"v"},
26+
Usage: "print additional data about each trip and vehicle",
27+
},
28+
&cli.StringFlag{
29+
Name: "extension",
30+
Usage: "GTFS realtime extension to use: nycttrips, nyctalerts",
31+
},
32+
},
2033
Commands: []*cli.Command{
2134
{
2235
Name: "static",
@@ -36,19 +49,8 @@ func main() {
3649
},
3750
},
3851
{
39-
Name: "realtime",
40-
Usage: "parse a GTFS realtime message",
41-
Flags: []cli.Flag{
42-
&cli.BoolFlag{
43-
Name: "verbose",
44-
Aliases: []string{"v"},
45-
Usage: "print additional data about each trip and vehicle",
46-
},
47-
&cli.StringFlag{
48-
Name: "extension",
49-
Usage: "GTFS realtime extension to use: nycttrips, nyctalerts",
50-
},
51-
},
52+
Name: "realtime",
53+
Usage: "parse a GTFS realtime message",
5254
ArgsUsage: "path",
5355
Action: func(ctx *cli.Context) error {
5456
args := ctx.Args()
@@ -60,27 +62,10 @@ func main() {
6062
if err != nil {
6163
return fmt.Errorf("failed to read file %s: %w", path, err)
6264
}
63-
6465
opts := gtfs.ParseRealtimeOptions{}
65-
switch ctx.String("extension") {
66-
case "nycttrips":
67-
opts.Extension = nycttrips.Extension(nycttrips.ExtensionOpts{
68-
FilterStaleUnassignedTrips: true,
69-
})
70-
americaNewYorkTimezone, err := time.LoadLocation("America/New_York")
71-
if err == nil {
72-
opts.Timezone = americaNewYorkTimezone
73-
}
74-
case "nyctalerts":
75-
opts.Extension = nyctalerts.Extension(nyctalerts.ExtensionOpts{
76-
ElevatorAlertsDeduplicationPolicy: nyctalerts.DeduplicateInComplex,
77-
ElevatorAlertsInformUsingStationIDs: true,
78-
SkipTimetabledNoServiceAlerts: true,
79-
})
80-
americaNewYorkTimezone, err := time.LoadLocation("America/New_York")
81-
if err == nil {
82-
opts.Timezone = americaNewYorkTimezone
83-
}
66+
rawExtension := ctx.String("extension")
67+
if err := readGtfsRealtimeExtension(rawExtension, &opts); err != nil {
68+
return err
8469
}
8570
realtime, err := gtfs.ParseRealtime(b, &opts)
8671
if err != nil {
@@ -102,6 +87,65 @@ func main() {
10287
return nil
10388
},
10489
},
90+
{
91+
Name: "journal",
92+
Usage: "build a journal from a series of GTFS realtime messages",
93+
Flags: []cli.Flag{
94+
&cli.StringFlag{
95+
Name: "output",
96+
Aliases: []string{"o"},
97+
Usage: "directory to output the CSV files",
98+
},
99+
},
100+
ArgsUsage: "path",
101+
Action: func(ctx *cli.Context) error {
102+
args := ctx.Args()
103+
if args.Len() == 0 {
104+
return fmt.Errorf("a path to the GTFS realtime messages was not provided")
105+
}
106+
path := ctx.Args().First()
107+
opts := gtfs.ParseRealtimeOptions{}
108+
rawExtension := ctx.String("extension")
109+
if err := readGtfsRealtimeExtension(rawExtension, &opts); err != nil {
110+
return err
111+
}
112+
113+
source, err := journal.NewDirectoryGtfsrtSource(path)
114+
if err != nil {
115+
return fmt.Errorf("failed to open %s: %w", path, err)
116+
}
117+
fmt.Println("Building journal...")
118+
j := journal.BuildJournal(source, time.Unix(0, 0), time.Now())
119+
fmt.Println("Exporting journal to CSV format...")
120+
export, err := j.ExportToCsv()
121+
if err != nil {
122+
return fmt.Errorf("failed to export journal: %w", err)
123+
}
124+
125+
outputDir := ctx.String("output")
126+
for _, f := range []struct {
127+
file string
128+
data []byte
129+
}{
130+
{
131+
file: "trips.csv",
132+
data: export.TripsCsv,
133+
},
134+
{
135+
file: "stop_times.csv",
136+
data: export.StopTimesCsv,
137+
},
138+
} {
139+
fullPath := filepath.Join(outputDir, f.file)
140+
fmt.Printf("Writing %s to %s\n", f.file, fullPath)
141+
if err := os.WriteFile(fullPath, f.data, 0666); err != nil {
142+
return fmt.Errorf("failed to write %s: %w", f.file, err)
143+
}
144+
}
145+
fmt.Println("Done")
146+
return nil
147+
},
148+
},
105149
},
106150
}
107151
if err := app.Run(os.Args); err != nil {
@@ -110,6 +154,33 @@ func main() {
110154
}
111155
}
112156

157+
func readGtfsRealtimeExtension(s string, opts *gtfs.ParseRealtimeOptions) error {
158+
switch s {
159+
case "":
160+
case "nycttrips":
161+
opts.Extension = nycttrips.Extension(nycttrips.ExtensionOpts{
162+
FilterStaleUnassignedTrips: true,
163+
})
164+
americaNewYorkTimezone, err := time.LoadLocation("America/New_York")
165+
if err == nil {
166+
opts.Timezone = americaNewYorkTimezone
167+
}
168+
case "nyctalerts":
169+
opts.Extension = nyctalerts.Extension(nyctalerts.ExtensionOpts{
170+
ElevatorAlertsDeduplicationPolicy: nyctalerts.DeduplicateInComplex,
171+
ElevatorAlertsInformUsingStationIDs: true,
172+
SkipTimetabledNoServiceAlerts: true,
173+
})
174+
americaNewYorkTimezone, err := time.LoadLocation("America/New_York")
175+
if err == nil {
176+
opts.Timezone = americaNewYorkTimezone
177+
}
178+
default:
179+
return fmt.Errorf("unknown extension %q; supported extensions are nycttrips and nyctalerts", s)
180+
}
181+
return nil
182+
}
183+
113184
func formatTrip(trip gtfs.Trip, indent int, printStopTimes bool) string {
114185
var b strings.Builder
115186
tc := color.New(color.FgCyan)

‎journal/export.go‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,69 @@
1+
package journal
2+
3+
import (
4+
"bytes"
5+
_ "embed"
6+
"fmt"
7+
"text/template"
8+
"time"
9+
10+
"github.com/jamespfennell/gtfs"
11+
)
12+
13+
//go:embed trips.csv.tmpl
14+
var tripsCsvTmpl string
15+
16+
//go:embed stop_times.csv.tmpl
17+
var stopTimesCsvTmpl string
18+
19+
var funcMap = template.FuncMap{
20+
"NullableString": func(s *string) string {
21+
if s == nil {
22+
return ""
23+
}
24+
return *s
25+
},
26+
"NullableUnix": func(t *time.Time) string {
27+
if t == nil {
28+
return ""
29+
}
30+
return fmt.Sprintf("%d", t.Unix())
31+
},
32+
"FormatDirectionID": func(d gtfs.DirectionID) string {
33+
switch d {
34+
case gtfs.DirectionID_False:
35+
return "0"
36+
case gtfs.DirectionID_True:
37+
return "1"
38+
default:
39+
return ""
40+
}
41+
},
42+
}
43+
44+
var tripsCsv *template.Template = template.Must(template.New("trips.csv.tmpl").Funcs(funcMap).Parse(tripsCsvTmpl))
45+
var stopTimesCsv *template.Template = template.Must(template.New("stop_times.csv.tmpl").Funcs(funcMap).Parse(stopTimesCsvTmpl))
46+
47+
// CsvExport contains CSV exports of a journal
48+
type CsvExport struct {
49+
TripsCsv []byte
50+
StopTimesCsv []byte
51+
}
52+
53+
func (journal *Journal) ExportToCsv() (*CsvExport, error) {
54+
var tripsB bytes.Buffer
55+
err := tripsCsv.Execute(&tripsB, journal.Trips)
56+
if err != nil {
57+
return nil, err
58+
}
59+
60+
var stopTimesB bytes.Buffer
61+
err = stopTimesCsv.Execute(&stopTimesB, journal.Trips)
62+
if err != nil {
63+
return nil, err
64+
}
65+
return &CsvExport{
66+
TripsCsv: tripsB.Bytes(),
67+
StopTimesCsv: stopTimesB.Bytes(),
68+
}, nil
69+
}

‎journal/export_test.go‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
package journal
2+
3+
import (
4+
"testing"
5+
"time"
6+
7+
"github.com/jamespfennell/gtfs"
8+
)
9+
10+
var trip Trip = Trip{
11+
TripUID: "TripUID",
12+
TripID: "TripID",
13+
RouteID: "RouteID",
14+
DirectionID: gtfs.DirectionID_True,
15+
VehicleID: "VehicleID",
16+
StartTime: time.Unix(100, 0),
17+
StopTimes: []StopTime{
18+
{
19+
StopID: "StopID1",
20+
Track: ptr("Track1"),
21+
ArrivalTime: nil,
22+
DepartureTime: ptr(time.Unix(200, 0)),
23+
LastObserved: time.Unix(200, 0),
24+
MarkedPast: ptr(time.Unix(300, 0)),
25+
},
26+
{
27+
StopID: "StopID2",
28+
ArrivalTime: ptr(time.Unix(300, 0)),
29+
DepartureTime: ptr(time.Unix(400, 0)),
30+
LastObserved: time.Unix(400, 0),
31+
},
32+
{
33+
StopID: "StopID3",
34+
Track: ptr("Track3"),
35+
ArrivalTime: ptr(time.Unix(500, 0)),
36+
DepartureTime: nil,
37+
LastObserved: time.Unix(400, 0),
38+
},
39+
},
40+
LastObserved: time.Unix(400, 0),
41+
MarkedPast: ptr(time.Unix(600, 0)),
42+
NumUpdates: 100,
43+
NumScheduleChanges: 2,
44+
NumScheduleRewrites: 1,
45+
}
46+
47+
const expectedTripsCsv = `trip_uid,trip_id,route_id,direction_id,start_time,vehicle_id,last_observed,marked_past,num_updates,num_schedule_changes,num_schedule_rewrites
48+
TripUID,TripID,RouteID,1,100,VehicleID,400,600,100,2,1
49+
`
50+
51+
const expectedStopTimesCsv = `trip_uid,stop_id,track,arrival_time,departure_time,last_observed,marked_past
52+
TripUID,StopID1,Track1,,200,200,300
53+
TripUID,StopID2,,300,400,400,
54+
TripUID,StopID3,Track3,500,,400,
55+
`
56+
57+
func TestCsvExport(t *testing.T) {
58+
journal := Journal{Trips: []Trip{trip}}
59+
60+
result, err := journal.ExportToCsv()
61+
if err != nil {
62+
t.Fatalf("AsCsv function failed: %s", err)
63+
}
64+
65+
if got, want := string(result.TripsCsv), expectedTripsCsv; got != want {
66+
t.Errorf("Trips file actual:\n%s\n!= expected:\n%s\n", got, want)
67+
}
68+
69+
if got, want := string(result.StopTimesCsv), expectedStopTimesCsv; got != want {
70+
t.Errorf("Stop times file actual:\n%s\n!= expected:\n%s\n", got, want)
71+
}
72+
}

‎journal/journal_test.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -187,6 +187,6 @@ func mtp(r int64) *time.Time {
187187
return &t
188188
}
189189

190-
func ptr(t string) *string {
190+
func ptr[T any](t T) *T {
191191
return &t
192192
}

‎journal/stop_times.csv.tmpl‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
trip_uid,stop_id,track,arrival_time,departure_time,last_observed,marked_past
2+
{{ range $trip := . -}}
3+
{{- range .StopTimes -}}
4+
{{- $trip.TripUID }},{{ .StopID }},{{ NullableString .Track }},{{ NullableUnix .ArrivalTime }},{{ NullableUnix .DepartureTime }},{{ .LastObserved.Unix }},{{ NullableUnix .MarkedPast }}
5+
{{ end -}}
6+
{{ end -}}

‎journal/trips.csv.tmpl‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
trip_uid,trip_id,route_id,direction_id,start_time,vehicle_id,last_observed,marked_past,num_updates,num_schedule_changes,num_schedule_rewrites
2+
{{ range . -}}
3+
{{ .TripUID }},{{ .TripID }},{{ .RouteID }},{{ FormatDirectionID .DirectionID }},{{ .StartTime.Unix }},{{ .VehicleID }},{{ .LastObserved.Unix }},{{ NullableUnix .MarkedPast }},{{ .NumUpdates }},{{ .NumScheduleChanges }},{{ .NumScheduleRewrites }}
4+
{{ end -}}

0 commit comments

Comments
 (0)