diff --git a/dcc/ss.py b/dcc/ss.py index 2579a3b..6d5dc2e 100644 --- a/dcc/ss.py +++ b/dcc/ss.py @@ -5,32 +5,48 @@ import sys import wand.display import wand.image + class SS(dcc.doom_base.WadMap): - def get_parser(self, prog_name): - parser = super().get_parser(prog_name) - parser.add_argument("-g", "--gravity", default="center") - parser.add_argument("-y", "--yolo", action="store_true") - return parser + def get_parser(self, prog_name): + parser = super().get_parser(prog_name) + parser.add_argument("-g", "--gravity", default="center") + parser.add_argument("-y", "--yolo", action="store_true") + return parser - def take_action(self, parsed_args): - def try_screenshot(): - with wand.image.Image(width=dcc.config.THUMB_WIDTH, height=dcc.config.THUMB_HEIGHT, pseudo="x:") as img: - img.reset_coords() - if img.size[0] < dcc.config.THUMB_WIDTH or img.size[1] < dcc.config.THUMB_HEIGHT: - if not messagebox.askretrycancel(title="DCC", message="Image too small. Try again?"): - sys.exit("Gave up trying to select an image.") - return False - img.crop(width=dcc.config.THUMB_WIDTH,height=dcc.config.THUMB_HEIGHT,gravity=parsed_args.gravity) - img.reset_coords() - if not parsed_args.yolo: - wand.display.display(img) - accepted = messagebox.askyesnocancel(title="DCC", message="Is this image acceptable?") - if accepted is None: - sys.exit("Gave up on image verification") - if not accepted: - return False - img.save(filename=self.base_thumb_path()) - return True + def take_action(self, parsed_args): + while not self._try_screenshot(parsed_args.gravity, parsed_args.yolo): + pass - while not try_screenshot(): - pass + def _try_screenshot(self, gravity, yolo): + with wand.image.Image( + width=dcc.config.THUMB_WIDTH, + height=dcc.config.THUMB_HEIGHT, + pseudo="x:" + ) as img: + img.reset_coords() + if ( + img.size[0] < dcc.config.THUMB_WIDTH + or img.size[1] < dcc.config.THUMB_HEIGHT + ): + if not messagebox.askretrycancel( + title="DCC", message="Image too small. Try again?" + ): + sys.exit("Gave up trying to select an image.") + return False + img.crop( + width=dcc.config.THUMB_WIDTH, + height=dcc.config.THUMB_HEIGHT, + gravity=gravity + ) + img.reset_coords() + if not yolo: + wand.display.display(img) + accepted = messagebox.askyesnocancel( + title="DCC", message="Is this image acceptable?" + ) + if accepted is None: + sys.exit("Gave up on image verification") + if not accepted: + return False + img.save(filename=self.base_thumb_path()) + return True