1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
|
// gtfsdb is a command line tool to roundtrip GTFS data to MySQL..
// It's especially powerful when used with Dolt.
package main
import (
"archive/zip"
"bytes"
"fmt"
"io"
"net/http"
"os"
"path/filepath"
"strings"
"encoding/csv"
)
type GTFS map[string][]string
func readFromURL(url string) GTFS {
// Get GTFS zip file from URL
resp, err := http.Get(url)
if err != nil {
fmt.Println(err)
os.Exit(1)
}
defer resp.Body.Close()
// Read all into a byte slice, then make a ReaderAt from it
body, err := io.ReadAll(resp.Body)
if err != nil {
fmt.Println(err)
os.Exit(1)
}
reader := bytes.NewReader(body)
// Read zip file
return readFromZip(reader)
}
func readFromFile(filename string) GTFS {
// Open GTFS zip file
file, err := os.Open(filename)
if err != nil {
fmt.Println(err)
os.Exit(1)
}
defer file.Close()
// Read zip file
return readFromZip(file)
}
func readFromZip(file io.ReaderAt) GTFS {
// Open zip file
r, err := zip.NewReader(file, 0)
if err != nil {
fmt.Println(err)
os.Exit(1)
}
// Read zip file
gtfs := make(GTFS)
for _, f := range r.File {
// Open file
rc, err := f.Open()
if err != nil {
fmt.Println(err)
os.Exit(1)
}
defer rc.Close()
// Read file into GTFS, removing .txt extension
table := strings.TrimSuffix(f.Name, filepath.Ext(f.Name))
gtfs[table] = readFromCSV(rc)
}
return gtfs
}
func readFromCSV(file io.Reader) []string {
// Read CSV file
r := csv.NewReader(file)
rows, err := r.ReadAll()
if err != nil {
fmt.Println(err)
os.Exit(1)
}
// Convert rows to []string
var result []string
for _, row := range rows {
result = append(result, strings.Join(row, ","))
}
return result
}
func writeToDir(gtfs GTFS) {
for table, rows := range gtfs {
// Open file
file, err := os.Create(table + ".csv")
if err != nil {
fmt.Println(err)
os.Exit(1)
}
defer file.Close()
// Write rows
for _, row := range rows {
_, err := io.WriteString(file, row+"\n")
if err != nil {
fmt.Println(err)
os.Exit(1)
}
}
}
}
func main() {
// Read filename from command line and write contents of zip to current directory as CSV files
var filename string
if len(os.Args) > 1 {
filename = os.Args[1]
} else {
fmt.Println("Please provide a URL or filename")
os.Exit(1)
}
// Read GTFS zip file from URL or filename
var gtfs GTFS
if strings.HasPrefix(filename, "http") {
gtfs = readFromURL(filename)
} else {
gtfs = readFromFile(filename)
}
// Write GTFS to current directory
writeToDir(gtfs)
}
|
