From d2b2a54a724db89447027e44ecb8f9bbed72cb83 Mon Sep 17 00:00:00 2001
From: jvoisin <julien.voisin@dustri.org>
Date: Sun, 29 Apr 2018 22:55:26 +0200
Subject: [PATCH] MAT2's cli now uses meaningful return codes
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- Simplify the multiprocessing by using a Pool
- Use some functional (♥) constructions to exit
  with a return code
- Add some tests to prove that we're doing things
  that are working correctly
---
 main.py               | 42 ++++++++++++------------------------------
 tests/test_climat2.py | 16 ++++++++++++++++
 2 files changed, 28 insertions(+), 30 deletions(-)

diff --git a/main.py b/main.py
index 9e07e15..933059b 100755
--- a/main.py
+++ b/main.py
@@ -1,11 +1,12 @@
 #!/usr/bin/python3
 
 import os
+from typing import Tuple
+import sys
+import itertools
 import mimetypes
 import argparse
-from threading import Thread
 import multiprocessing
-from queue import Queue
 
 from src import parser_factory
 
@@ -52,14 +53,15 @@ def show_meta(filename:str):
             print("  %s: harmful content" % k)
 
 
-def clean_meta(filename:str, is_lightweigth:bool) -> bool:
+def clean_meta(params:Tuple[str, bool]) -> bool:
+    filename, is_lightweigth = params
     if not __check_file(filename, os.R_OK|os.W_OK):
         return
 
     p, mtype = parser_factory.get_parser(filename)
     if p is None:
         print("[-] %s's format (%s) is not supported" % (filename, mtype))
-        return
+        return False
     if is_lightweigth:
         return p.remove_all_lightweight()
     return p.remove_all()
@@ -82,15 +84,6 @@ def __get_files_recursively(files):
                 for _f in _files:
                     yield os.path.join(path, _f)
 
-def __do_clean_async(is_lightweigth, q):
-    while True:
-        f = q.get()
-        if f is None:  # nothing more to process
-            return
-        clean_meta(f, is_lightweigth)
-        q.task_done()
-
-
 def main():
     arg_parser = create_arg_parser()
     args = arg_parser.parse_args()
@@ -106,24 +99,13 @@ def main():
             show_meta(f)
         return
 
-    else:  # Thread the cleaning
+    else:
+        p = multiprocessing.Pool()
         mode = (args.lightweight is True)
-        q = Queue(maxsize=0)
-        threads = list()
-        for f in __get_files_recursively(args.files):
-            q.put(f)
-
-        for _ in range(multiprocessing.cpu_count()):
-            worker = Thread(target=__do_clean_async, args=(mode, q))
-            worker.start()
-            threads.append(worker)
-
-        for _ in range(multiprocessing.cpu_count()):
-            q.put(None)  # stop the threads
-
-        for worker in threads:
-            worker.join()
+        l = zip(__get_files_recursively(args.files), itertools.repeat(mode))
 
+        ret = list(p.imap_unordered(clean_meta, list(l)))
+        return 0 if all(ret) else -1
 
 if __name__ == '__main__':
-    main()
+    sys.exit(main())
diff --git a/tests/test_climat2.py b/tests/test_climat2.py
index 0536646..bc4a175 100644
--- a/tests/test_climat2.py
+++ b/tests/test_climat2.py
@@ -16,6 +16,22 @@ class TestHelp(unittest.TestCase):
         self.assertIn(b'usage: main.py [-h] [-c] [-l] [-s] [-L] [files [files ...]]', stdout)
 
 
+class TestReturnValue(unittest.TestCase):
+    def test_nonzero(self):
+        ret = subprocess.call(['./main.py', './main.py'], stdout=subprocess.DEVNULL)
+        self.assertEqual(255, ret)
+
+        ret = subprocess.call(['./main.py', '--whololo'], stderr=subprocess.DEVNULL)
+        self.assertEqual(2, ret)
+
+    def test_zero(self):
+        ret = subprocess.call(['./main.py'], stdout=subprocess.DEVNULL)
+        self.assertEqual(0, ret)
+
+        ret = subprocess.call(['./main.py', '--show', './main.py'], stdout=subprocess.DEVNULL)
+        self.assertEqual(0, ret)
+
+
 class TestCleanMeta(unittest.TestCase):
     def test_jpg(self):
         shutil.copy('./tests/data/dirty.jpg', './tests/data/clean.jpg')