From 34d7ee3c27feba1c45f605f1543864cada9a1b80 Mon Sep 17 00:00:00 2001
From: Alex Eidt <alex.eidt@outlook.com>
Date: Sat, 23 Apr 2022 20:19:15 -0700
Subject: [PATCH] Cleanup up imageio

---
 imageio.go | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/imageio.go b/imageio.go
index 167ce71..9b09cb6 100644
--- a/imageio.go
+++ b/imageio.go
@@ -5,7 +5,7 @@ import (
 	"fmt"
 	"image"
 	"os"
-	"strings"
+	"path/filepath"
 
 	"image/color"
 	"image/jpeg"
@@ -75,14 +75,16 @@ func Write(filename string, width, height int, buffer []byte) error {
 		}
 	}
 
-	if strings.HasSuffix(filename, ".png") {
+	switch filepath.Ext(filename) {
+	case ".png":
 		if err := png.Encode(f, image); err != nil {
 			return err
 		}
-	} else if strings.HasSuffix(filename, ".jpg") || strings.HasSuffix(filename, ".jpeg") {
+	case ".jpg", ".jpeg":
 		if err := jpeg.Encode(f, image, nil); err != nil {
 			return err
 		}
 	}
+
 	return nil
 }