From 2fb48c494b9e851e1fcbd660bab088118bb95bb9 Mon Sep 17 00:00:00 2001
From: Remy Moll <me@moll.re>
Date: Wed, 8 Nov 2023 22:26:01 +0100
Subject: [PATCH] cleaner test logic

---
 image_convert.py | 14 ++++++++------
 image_show.py    | 34 +++++++++++++++++++++++++++++++++-
 main.py          | 35 +++++++++++++++++++++--------------
 3 files changed, 62 insertions(+), 21 deletions(-)

diff --git a/image_convert.py b/image_convert.py
index 0829440..89fb564 100644
--- a/image_convert.py
+++ b/image_convert.py
@@ -1,6 +1,5 @@
 from PIL import Image, ImageOps
 import io
-import os
 
 black_base = [0, 0, 0]
 white_base = [255, 255, 255]
@@ -33,16 +32,16 @@ class ImageShrink:
     """Shrinks a given image (bytearray) to a given resolution (width, height)"""
     resolution = (480, 800)
 
-    def __init__(self) -> None:
-        pass
+    def __init__(self, dither=True, colors=7) -> None:
+        self.dither = dither
+        self.colors = colors
+
 
     def convert(self, image: bytearray) -> Image:
         # load image from bytearray
         image = Image.open(io.BytesIO(image))
         image = self.shrink(image)
         image = self.convert_to_reduced_colors(image)
-        if os.uname().machine == "x86_64":
-            image.save("test.png")
         return image
 
 
@@ -62,5 +61,8 @@ class ImageShrink:
         if image.mode != "RGB":
             print("Converting image to RGB")
             image = image.convert("RGB")
-        new_image = image.quantize(colors = len(palette), palette=ref_image, dither=True)
+        if self.colors == -1:
+            print("Skipping color reduction")
+            return image
+        new_image = image.quantize(colors = len(palette), palette=ref_image, dither=self.dither)
         return new_image
diff --git a/image_show.py b/image_show.py
index f154a3e..f056176 100644
--- a/image_show.py
+++ b/image_show.py
@@ -1,5 +1,37 @@
 from PIL import Image, ImageDraw
-import epaper
+try:
+    import epaper
+except ModuleNotFoundError:
+    # generate dummy class for testing
+    class epaper:
+        class epaper:
+            def __init__(self, name):
+                pass
+            class EPD:
+                width = 480
+                height = 800
+                WHITE = 0
+                BLUE = 0
+                ORANGE = 0
+                BLACK = 0
+                GREEN = 0
+                RED = 0
+                YELLOW = 0
+                def __init__(self):
+                    pass
+                def init(self):
+                    pass
+                def Clear(self):
+                    pass
+                def display(self, image):
+                    print("Dummy output")
+                def sleep(self):
+                    pass
+                def getbuffer(self, image:Image):
+                    image.save("test.png")
+                    return image
+                
+
 # BLACK  = 0x000000   #   0000  BGR
 # WHITE  = 0xffffff   #   0001
 # GREEN  = 0x00ff00   #   0010
diff --git a/main.py b/main.py
index c16a712..b73347c 100644
--- a/main.py
+++ b/main.py
@@ -1,20 +1,27 @@
+import sys
 from image_convert import ImageShrink
 from image_get import ImageGetter
-try:
-    from image_show import ImageShow
-    show_image = True
-except ImportError:
-    print("ImageShow not found. Image will not be displayed.")
-    show_image = False
-  
-get = ImageGetter()
-convert = ImageShrink()
-if show_image:
+from image_show import ImageShow
+
+if len(sys.argv) == 2 and sys.argv[1] == "test":
+    print("Running test")
     show = ImageShow()
+    show.draw_sample_image()
+    sys.exit()
+
+
+shrink_kwargs = {}
+if "nodither" in sys.argv:
+    print("Disabling dithering")
+    shrink_kwargs["dither"] = False
+if "noreduce" in sys.argv:
+    print("Disabling color reduction")
+    shrink_kwargs["colors"] = -1
+
+get = ImageGetter()
+convert = ImageShrink(**shrink_kwargs)
+show = ImageShow()
 
 image = get.get_random_image()
 image = convert.convert(image)
-if show_image:
-    show.show_image(image)
-else:
-    print("Would have shown image now.")
+show.show_image(image)