refactor: Moved multithreading to doctr.utils (#371)

* refactor: Moved multithreading to doctr.utils * refactor: Reflected changes on imports * test: Added unittest for multithreading * test: Removed legacy multithreading test * test: Fixed unittest for single threading
mindee · Jul 8, 2021 · b97363c · b97363c
1 parent 1910ed3
commit b97363c
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 16 deletions.
diff --git a/doctr/datasets/loader.py b/doctr/datasets/loader.py
@@ -8,7 +8,7 @@
 import numpy as np
 from typing import Optional
 
-from .multithreading import multithread_exec
+from doctr.utils.multithreading import multithread_exec
 
 __all__ = ["DataLoader"]
 

diff --git a/doctr/datasets/multithreading.py → doctr/utils/multithreading.py b/doctr/datasets/multithreading.py → doctr/utils/multithreading.py
diff --git a/test/common/test_utils_multithreading.py b/test/common/test_utils_multithreading.py
@@ -0,0 +1,21 @@
+import pytest
+
+import numpy as np
+from doctr.utils.multithreading import multithread_exec
+
+
+@pytest.mark.parametrize(
+    "input_seq, func, output_seq",
+    [
+        [[1, 2, 3], lambda x: 2 * x, [2, 4, 6]],
+        [[1, 2, 3], lambda x: x ** 2, [1, 4, 9]],
+        [
+            ['this is', 'show me', 'I know'],
+            lambda x: x + ' the way',
+            ['this is the way', 'show me the way', 'I know the way']
+        ],
+    ],
+)
+def test_multithread_exec(input_seq, func, output_seq):
+    assert multithread_exec(func, input_seq) == output_seq
+    assert list(multithread_exec(func, input_seq, 0)) == output_seq
diff --git a/test/tensorflow/test_datasets_multithreading_tf.py b/test/tensorflow/test_datasets_multithreading_tf.py