diff options
-rw-r--r-- | main.go | 36 |
1 files changed, 20 insertions, 16 deletions
@@ -9,7 +9,6 @@ import ( "io" "net/http" "os" - "path/filepath" "strings" "encoding/csv" ) @@ -25,20 +24,19 @@ func readFromURL(url string) GTFS { } defer resp.Body.Close() - // Read all into a byte slice, then make a ReaderAt from it + // Read all into a byte slice 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) + return readFromZip(body) } func readFromFile(filename string) GTFS { - // Open GTFS zip file + // Read GTFS zip file from filename file, err := os.Open(filename) if err != nil { fmt.Println(err) @@ -46,20 +44,27 @@ func readFromFile(filename string) GTFS { } defer file.Close() + // Read all into a byte slice + fbytes, err := io.ReadAll(file) + if err != nil { + fmt.Println(err) + os.Exit(1) + } + // Read zip file - return readFromZip(file) + return readFromZip(fbytes) } -func readFromZip(file io.ReaderAt) GTFS { +func readFromZip(file []byte) GTFS { // Open zip file - r, err := zip.NewReader(file, 0) + r, err := zip.NewReader(bytes.NewReader(file), int64(len(file))) if err != nil { fmt.Println(err) os.Exit(1) } - // Read zip file - gtfs := make(GTFS) + // Read all files in zip + gtfs := GTFS{} for _, f := range r.File { // Open file rc, err := f.Open() @@ -69,9 +74,8 @@ func readFromZip(file io.ReaderAt) GTFS { } defer rc.Close() - // Read file into GTFS, removing .txt extension - table := strings.TrimSuffix(f.Name, filepath.Ext(f.Name)) - gtfs[table] = readFromCSV(rc) + // Read CSV file. Remove .txt extension. + gtfs[f.Name[:len(f.Name)-4]] = readFromCSV(rc) } return gtfs @@ -95,10 +99,10 @@ func readFromCSV(file io.Reader) []string { return result } -func writeToDir(gtfs GTFS) { +func writeToDir(gtfs GTFS, ext string) { for table, rows := range gtfs { // Open file - file, err := os.Create(table + ".csv") + file, err := os.Create(table + "." + ext) if err != nil { fmt.Println(err) os.Exit(1) @@ -135,5 +139,5 @@ func main() { } // Write GTFS to current directory - writeToDir(gtfs) + writeToDir(gtfs, "csv") } |