add --threads option

diff --git a/main.py b/main.py
index bb3a40985c53c8b07e3937aa19c57c19aaa674e2..1d5918199d8b8fa3da7d7501b02af7e0791135c0 100755 (executable)
--- a/main.py
+++ b/main.py
@@ -8,9 +8,6 @@ logging.basicConfig(filename='/var/log/mqlg.log',
                                        format='[mq-lg] %(threadName)s %(asctime)s,%(msecs)d %(name)s %(levelname)s %(message)s',
                                        datefmt='%H:%M:%S',
                                        level=logging.INFO)
-printed = logging.StreamHandler()
-printed.setLevel(logging.INFO)#logging.WARNING)
-#logging.getLogger().addHandler(printed)
 
 # Queues
 FF_TAR="ff_tar"
@@ -175,17 +172,60 @@ def do_listen(message):
             f.write(content)
         logging.warning("saved {}".format(output_path))
 
+class InvalidArguments(Exception):
+    pass
 if __name__ == '__main__':
-    args = sys.argv
-    if   len(args) == 3 and sys.argv[1] == "sm" and sys.argv[2] == "worker":
-        threaded_consumer(sm.QUEUE, sm.QUEUE_RESP, functools.partial(do_work, sm.extract_text))
-    elif len(args) == 3 and sys.argv[1] == "sm" and sys.argv[2] == "listener":
-        threaded_consumer(sm.QUEUE_RESP, do_listen)
-    elif len(args) == 3 and sys.argv[1] == "sm" and sys.argv[2] == "enqueue":
-        enqueue_all(sm.QUEUE, sm.generate_tasks())
-    elif len(args) == 5 and sys.argv[1] == "sm" and sys.argv[2] == "debug":
-        printed.setLevel(logging.DEBUG)
-        sm.extract_text(args[3], args[4], debug=True)
-    else:
-        print("Usage: {} sm worker|listener|enqueue".format(sys.argv[0].split("/")[-1]))
+    args = sys.argv[1:]
+
+    # Options
+    threads  = None
+
+    try:
+        # Parse flag options
+        new_args = []
+        while len(args) > 0:
+            arg, *args = args
+            if arg == "--debug":
+                printed = logging.StreamHandler()
+                printed.setLevel(logging.INFO)
+                logging.getLogger().addHandler(printed)
+            elif arg == "--threads":
+                if len(args) < 1:
+                    raise InvalidArguments()
+                threads, args = int(args[0]), args[1:]
+            elif arg.startswith("--"):
+                raise InvalidArguments()
+            else:
+                new_args.append(arg)
+
+        # Parse command + dataset, two mandatory options
+        if len(new_args) < 2:
+            raise InvalidArguments()
+        dataset, command, *args = new_args
+
+        if dataset == "sm":
+            module = sm
+        else:
+            raise InvalidArguments()
+
+        if command == "worker":
+            if len(args) != 0:
+                raise InvalidArguments()
+            threaded_consumer(module.QUEUE, module.QUEUE_RESP, functools.partial(do_work, module.extract_text), threads=threads)
+        elif command == "listener":
+            if len(args) != 0:
+                raise InvalidArguments()
+            threaded_consumer(module.QUEUE_RESP, do_listen, threads=threads)
+        elif command == "enqueue":
+            if len(args) != 0:
+                raise InvalidArguments()
+            enqueue_all(queue, module.generate_tasks())
+        elif command == "debug":
+            if len(args) != 2:
+                raise InvalidArguments()
+            printed.setLevel(logging.DEBUG)
+            module.extract_text(*args, debug=True)
+    except InvalidArguments:
+        program_name = sys.argv[0].split("/")[-1]
+        print("Usage: {} sm worker|listener|enqueue".format(program_name))
         sys.exit(1)