Finished exercise and improved error handling

star-shapes/main.go

@@ -1,32 +1,42 @@
 package main
 
-import "fmt"
+// TODO: figure out colored text
+
+import (
+	"fmt"
+)
 
 func main() {
 	var numStar int
 	var numShape int
 
+	// TODO: Add error handling for star number
 	fmt.Println("How many stars do you want?")
 	fmt.Printf("    Stars: ")
 	fmt.Scan(&numStar)
 
 	fmt.Printf("\n")
 
+forShapes:
+	for {
 		fmt.Println("Which shape would you like?\n    Line [1]\n    Square [2]\n    Flag [3]")
 		fmt.Printf("Shape: ")
-	fmt.Scan(&numShape)
+		fmt.Scanln(&numShape)
-
+		fmt.Printf("\n")
 		switch numShape {
 		case 1:
-		println("Line")
+			line(numStar)
+			break forShapes
 		case 2:
-		fmt.Println("Square")
+			square(numStar)
+			break forShapes
 		case 3:
-		fmt.Println("Flag")
+			flag(numStar)
+			break forShapes
 		default:
-		fmt.Println("Invalid input, please use either 1, 2 or 3.")
+			fmt.Printf("\n[-] INVALID INPUT, please use either 1, 2 or 3.\n\n")
+		}
 	}
-	//flag(numStar)
 }
 
 func line(num int) {