From 47e787c80c7292dd2c866a4a9065fedf74b845ee Mon Sep 17 00:00:00 2001
From: k4yt3x <k4yt3x@k4yt3x.com>
Date: Sat, 27 Jul 2019 17:52:31 -0400
Subject: [PATCH] added more exceptions for better error handling

---
 bin/exceptions.py | 15 +++++++++++++++
 bin/upscaler.py   |  8 ++++----
 bin/video2x.py    | 13 +++++++------
 3 files changed, 26 insertions(+), 10 deletions(-)

diff --git a/bin/exceptions.py b/bin/exceptions.py
index 6e8029a..ad6346d 100644
--- a/bin/exceptions.py
+++ b/bin/exceptions.py
@@ -11,3 +11,18 @@ Last Modified: July 27, 2019
 class ArgumentError(Exception):
     def __init__(self, message):
         super().__init__(message)
+
+
+class StreamNotFoundError(Exception):
+    def __init__(self, message):
+        super().__init__(message)
+
+
+class UnrecognizedDriverError(Exception):
+    def __init__(self, message):
+        super().__init__(message)
+
+
+class UnsupportedPixelError(Exception):
+    def __init__(self, message):
+        super().__init__(message)
diff --git a/bin/upscaler.py b/bin/upscaler.py
index 3e1faae..43d0902 100644
--- a/bin/upscaler.py
+++ b/bin/upscaler.py
@@ -157,7 +157,7 @@ class Upscaler:
         # initialize waifu2x driver
         drivers = ['waifu2x_caffe', 'waifu2x_converter', 'waifu2x_ncnn_vulkan']
         if self.waifu2x_driver not in drivers:
-            raise Exception(f'Unrecognized waifu2x driver: {self.waifu2x_driver}')
+            raise UnrecognizedDriverError(f'Unrecognized waifu2x driver: {self.waifu2x_driver}')
 
         # it's easier to do multi-threading with waifu2x_converter
         # the number of threads can be passed directly to waifu2x_converter
@@ -314,7 +314,7 @@ class Upscaler:
         # exit if no video stream found
         if video_stream_index is None:
             Avalon.error('Aborting: No video stream found')
-            exit(1)
+            raise StreamNotFoundError('no video stream found')
 
         # get average frame rate of video stream
         framerate = float(Fraction(video_info['streams'][video_stream_index]['avg_frame_rate']))
@@ -325,9 +325,9 @@ class Upscaler:
 
         try:
             self.bit_depth = pixel_formats[fm.pixel_format]
-        except KeyError as e:
+        except KeyError:
             Avalon.error(f'Unsupported pixel format: {fm.pixel_format}')
-            raise e
+            raise UnsupportedPixelError(f'unsupported pixel format {fm.pixel_format}')
 
         Avalon.info(f'Framerate: {framerate}')
 
diff --git a/bin/video2x.py b/bin/video2x.py
index 9ad8764..26be60d 100644
--- a/bin/video2x.py
+++ b/bin/video2x.py
@@ -43,6 +43,7 @@ smooth and edges sharp.
 """
 
 # local imports
+from exceptions import *
 from upscaler import Upscaler
 
 # built-in imports
@@ -254,22 +255,22 @@ if args.version:
 # arguments sanity check
 if not args.input:
     Avalon.error('You must specify input video file/directory path')
-    exit(1)
+    raise ArgumentError('input video path not specified')
 if not args.output:
     Avalon.error('You must specify output video file/directory path')
-    exit(1)
+    raise ArgumentError('output video path not specified')
 if (args.driver == 'waifu2x_converter' or args.driver == 'waifu2x_ncnn_vulkan') and args.width and args.height:
     Avalon.error('Waifu2x Converter CPP/NCNN accepts only scaling ratio')
-    exit(1)
+    raise ArgumentError('waifu2x-converter supports only scaling ratio')
 if args.driver == 'waifu2x_ncnn_vulkan' and (args.ratio > 2 or not args.ratio.is_integer()):
     Avalon.error('Scaling ratio must be 1 or 2 for waifu2x_ncnn_vulkan')
-    exit(1)
+    raise ArgumentError('scaling ratio must be 1 or 2 for waifu2x_ncnn_vulkan')
 if (args.width or args.height) and args.ratio:
     Avalon.error('You can only specify either scaling ratio or output width and height')
-    exit(1)
+    raise ArgumentError('both scaling ration and width/height specified')
 if (args.width and not args.height) or (not args.width and args.height):
     Avalon.error('You must specify both width and height')
-    exit(1)
+    raise ArgumentError('only one of width or height is specified')
 
 # check available memory
 check_memory()