Initial working Go codeHEADmain

1 files changed, 20 insertions, 16 deletions

diff --git a/main.go b/main.go
index 59c58c8..647ad2b 100644
--- a/main.go
+++ b/main.go
@@ -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")
 }