From e2d46253254c26fb1153e66df6c23bbcba375243 Mon Sep 17 00:00:00 2001
From: Eitan Simler <esimler@proofpoint.com>
Date: Tue, 4 Aug 2020 21:43:48 -0700
Subject: [PATCH] Use argparse choices param

---
 fawkes/protection.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/fawkes/protection.py b/fawkes/protection.py
index 108698e..2e01fa6 100644
--- a/fawkes/protection.py
+++ b/fawkes/protection.py
@@ -194,13 +194,12 @@ def main(*argv):
                         action='store_true')
     parser.add_argument('--debug', help="turn on debug and copy/paste the stdout when reporting an issue on github",
                         action='store_true')
-    parser.add_argument('--format', type=str,
+    parser.add_argument('--format', type=str, choices=('png', 'jpg', 'jpeg'),
                         help="format of the output image",
                         default="png")
 
     args = parser.parse_args(argv[1:])
 
-    assert args.format in ['png', 'jpg', 'jpeg']
     if args.format == 'jpg':
         args.format = 'jpeg'