Changed return type of multithread_exec to iterator

doctr/datasets/loader.py

@@ -90,7 +90,7 @@ def __next__(self):
             idx = self._num_yielded * self.batch_size
             indices = self.indices[idx: min(len(self.dataset), idx + self.batch_size)]
 
-            samples = multithread_exec(self.dataset.__getitem__, indices, threads=self.num_workers)
+            samples = list(multithread_exec(self.dataset.__getitem__, indices, threads=self.num_workers))
 
             batch_data = self.collate_fn(samples)
 

doctr/models/preprocessor/pytorch.py

@@ -117,13 +117,13 @@ def __call__(
 
         elif isinstance(x, list) and all(isinstance(sample, (np.ndarray, torch.Tensor)) for sample in x):
             # Sample transform (to tensor, resize)
-            samples = multithread_exec(self.sample_transforms, x)
+            samples = list(multithread_exec(self.sample_transforms, x))
             # Batching
-            batches = self.batch_inputs(samples)  # type: ignore[arg-type]
+            batches = self.batch_inputs(samples)
         else:
             raise TypeError(f"invalid input type: {type(x)}")
 
         # Batch transforms (normalize)
-        batches = multithread_exec(self.normalize, batches)  # type: ignore[assignment]
+        batches = list(multithread_exec(self.normalize, batches))
 
         return batches

doctr/models/preprocessor/tensorflow.py

@@ -115,13 +115,13 @@ def __call__(
 
         elif isinstance(x, list) and all(isinstance(sample, (np.ndarray, tf.Tensor)) for sample in x):
             # Sample transform (to tensor, resize)
-            samples = multithread_exec(self.sample_transforms, x)
+            samples = list(multithread_exec(self.sample_transforms, x))
             # Batching
-            batches = self.batch_inputs(samples)  # type: ignore[arg-type]
+            batches = self.batch_inputs(samples)
         else:
             raise TypeError(f"invalid input type: {type(x)}")
 
         # Batch transforms (normalize)
-        batches = multithread_exec(self.normalize, batches)  # type: ignore[assignment]
+        batches = list(multithread_exec(self.normalize, batches))
 
         return batches

doctr/utils/multithreading.py

@@ -6,12 +6,12 @@
 
 import multiprocessing as mp
 from multiprocessing.pool import ThreadPool
-from typing import Any, Callable, Iterable, Optional
+from typing import Any, Callable, Iterable, Iterator, Optional
 
 __all__ = ['multithread_exec']
 
 
-def multithread_exec(func: Callable[[Any], Any], seq: Iterable[Any], threads: Optional[int] = None) -> Iterable[Any]:
+def multithread_exec(func: Callable[[Any], Any], seq: Iterable[Any], threads: Optional[int] = None) -> Iterator[Any]:
     """Execute a given function in parallel for each element of a given sequence
 
     >>> from doctr.utils.multithreading import multithread_exec
@@ -24,7 +24,7 @@ def multithread_exec(func: Callable[[Any], Any], seq: Iterable[Any], threads: Op
         threads: number of workers to be used for multiprocessing
 
     Returns:
-        iterable of the function's results using the iterable as inputs
+        iterator of the function's results using the iterable as inputs
     """
 
     threads = threads if isinstance(threads, int) else min(16, mp.cpu_count())
@@ -34,5 +34,7 @@ def multithread_exec(func: Callable[[Any], Any], seq: Iterable[Any], threads: Op
     # Multi-threading
     else:
         with ThreadPool(threads) as tp:
-            results = tp.map(func, seq)  # type: ignore[assignment]
+            # ThreadPool's map function returns a list, but seq could be of a different type
+            # That's why wrapping result in map to return iterator
+            results = map(lambda x: x, tp.map(func, seq))
     return results

tests/common/test_utils_multithreading.py

@@ -16,5 +16,5 @@
     ],
 )
 def test_multithread_exec(input_seq, func, output_seq):
-    assert multithread_exec(func, input_seq) == output_seq
+    assert list(multithread_exec(func, input_seq)) == output_seq
     assert list(multithread_exec(func, input_seq, 0)) == output_seq