Initial working Go code

main
Joop Kiefte 3 weeks ago
parent b47b66288f
commit 73cb298e94

@ -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")
}

Loading…
Cancel
Save