improved argument reading

meme_otron/__main__.py

@@ -4,6 +4,7 @@ import os
 from . import img_factory
 from . import meme_db
 from . import meme_otron
+from . import utils
 from . import VERSION
 
 if __name__ == "__main__":
@@ -12,7 +13,7 @@ if __name__ == "__main__":
 
     # TODO better arguments reading (-h, -o, -v)
 
-    if len(sys.argv) <= 1 or sys.argv[1].lower().strip() == "help" or "-h" in sys.argv:
+    if len(sys.argv) <= 1 or utils.read_argument(sys.argv, "help", "--help", "-h"):
         print(f"Meme-Otron v{VERSION}"
               "python -m meme_otron -h\n"
               "python -m meme_otron (meme_id) \"[text 1]\" \"[text 2]\" ... > file.jpg\n"
@@ -20,13 +21,7 @@ if __name__ == "__main__":
               file=sys.stderr)
         sys.exit(1)
     else:
-        output_file = None
+        output_file = utils.read_argument(sys.argv, "-o", "--output", valued=True, delete=True)
-        if "-o" in sys.argv:
-            i = sys.argv.index("-o")
-            if len(sys.argv) >= i:
-                output_file = sys.argv[i + 1]
-                del sys.argv[i + 1]
-            del sys.argv[i]
         img = meme_otron.compute(*sys.argv[1:])
         if img is None:
             proposal = meme_db.find_nearest(sys.argv[1])

meme_otron/utils.py

@@ -143,3 +143,21 @@ def safe_index(src: Union[str, list], pattern, start: int = 0):
         return src.index(pattern, start)
     except ValueError:
         return None
+
+
+def read_argument(args: List[str], *names: str, valued: bool = False, delete: bool = False):
+    for i, arg in enumerate(args):
+        if arg.lower() in names:
+            if delete:
+                del args[i]
+                i -= 1
+            if not valued:
+                return True
+            else:
+                v = None
+                if i + 1 < len(args):
+                    v = args[i + 1]
+                    if delete:
+                        del args[i + 1]
+                return v
+    return None

tests/unit/meme_otron/test_utils.py

@@ -128,3 +128,18 @@ class Test(TestCase):
         self.assertEqual([5, 9, 15], utils.place_line_breaks([5.2, 14.3, 15.2], [3, 5, 9, 15, 18]))
         self.assertEqual([5, 9, 15, 18], utils.place_line_breaks([5.2, 14.3, 14.5, 15.2], [3, 5, 9, 15, 18]))
         self.assertEqual([5, 9, 15, 18], utils.place_line_breaks([5.2, 14.3, 14.5, 15.2], [3, 5, 9, 15, 18, 20]))
+
+    def test_read_argument(self):
+        self.assertIsNone(utils.read_argument(["test", "-o", "test"], "--output"))
+        self.assertTrue(utils.read_argument(["test", "-O", "test"], "--output", "-o"))
+        self.assertIsNone(utils.read_argument(["test", "-o"], "-o", valued=True))
+        self.assertEqual("test1", utils.read_argument(["test", "-o", "test1", "-o", "test2"], "-o", valued=True))
+        args = ["test", "-o", "test1"]
+        self.assertTrue(utils.read_argument(args, "-o", delete=True))
+        self.assertEqual(["test", "test1"], args)
+        args = ["test", "-o", "test1"]
+        self.assertEqual("test1", utils.read_argument(args, "-o", valued=True, delete=True))
+        self.assertEqual(["test"], args)
+        args = ["test", "-o"]
+        self.assertIsNone(utils.read_argument(args, "-o", valued=True, delete=True))
+        self.assertEqual(["test"], args)