restrict pillow to png

1 files changed, 2 insertions, 2 deletions

diff --git a/wfs.py b/wfs.py
index 0747bec..c107c67 100755
--- a/wfs.py
+++ b/wfs.py
@@ -149,14 +149,14 @@ def verify(path, screenshot_begin_png):
 
     from io import BytesIO
     from PIL import Image, ImageChops
-    screenshot_begin = Image.open(BytesIO(screenshot_begin_png)).convert('RGB')
+    screenshot_begin = Image.open(BytesIO(screenshot_begin_png), formats=('PNG',)).convert('RGB')
 
     logger.info('taking post-screenshot for %s', path)
     DRIVER.set_window_size(*screenshot_begin.size)
     screenshot_end_png = DRIVER.get_screenshot_as_png()
 
     logger.info('checking screenshot for %s', path)
-    screenshot_end = Image.open(BytesIO(screenshot_end_png)).convert('RGB')
+    screenshot_end = Image.open(BytesIO(screenshot_end_png), formats=('PNG',)).convert('RGB')
     if ImageChops.difference(screenshot_begin, screenshot_end).getbbox():
         raise Exception(f'screenshots do not match for {path}')