From 286ede5fd34ba44bbc71a3109b08d8519460ea89 Mon Sep 17 00:00:00 2001
From: Andras Schmelczer <andras@schmelczer.dev>
Date: Sat, 9 Jul 2022 21:33:14 +0200
Subject: [PATCH] Add simpler parallel_map

---
 great_ai/utilities/__init__.py                |  7 ++++-
 great_ai/utilities/parallel_map/__init__.py   |  1 +
 .../parallel_map/simple_parallel_map.py       | 30 +++++++++++++++++++
 3 files changed, 37 insertions(+), 1 deletion(-)
 create mode 100644 great_ai/utilities/parallel_map/simple_parallel_map.py

diff --git a/great_ai/utilities/__init__.py b/great_ai/utilities/__init__.py
index 5bc97c4..65c2a17 100644
--- a/great_ai/utilities/__init__.py
+++ b/great_ai/utilities/__init__.py
@@ -5,6 +5,11 @@ from .evaluate_ranking import evaluate_ranking
 from .get_sentences import get_sentences
 from .language import english_name_of_language, is_english, predict_language
 from .logger import get_logger
-from .parallel_map import WorkerException, parallel_map, threaded_parallel_map
+from .parallel_map import (
+    WorkerException,
+    parallel_map,
+    simple_parallel_map,
+    threaded_parallel_map,
+)
 from .unchunk import unchunk
 from .unique import unique
diff --git a/great_ai/utilities/parallel_map/__init__.py b/great_ai/utilities/parallel_map/__init__.py
index 7b6ceec..a984057 100644
--- a/great_ai/utilities/parallel_map/__init__.py
+++ b/great_ai/utilities/parallel_map/__init__.py
@@ -1,3 +1,4 @@
 from .parallel_map import parallel_map
+from .simple_parallel_map import simple_parallel_map
 from .threaded_parallel_map import threaded_parallel_map
 from .worker_exception import WorkerException
diff --git a/great_ai/utilities/parallel_map/simple_parallel_map.py b/great_ai/utilities/parallel_map/simple_parallel_map.py
new file mode 100644
index 0000000..3dca6fd
--- /dev/null
+++ b/great_ai/utilities/parallel_map/simple_parallel_map.py
@@ -0,0 +1,30 @@
+from typing import Awaitable, Callable, List, Optional, Sequence, TypeVar, Union
+
+from tqdm.cli import tqdm
+
+from .parallel_map import parallel_map
+
+T = TypeVar("T")
+V = TypeVar("V")
+
+
+def simple_parallel_map(
+    func: Callable[[T], Union[V, Awaitable[V]]],
+    input_values: Sequence[T],
+    *,
+    chunk_size: Optional[int] = None,
+    concurrency: Optional[int] = None,
+) -> List[V]:
+    input_values = list(input_values)  # in case the input is mistakenly not a sequence
+    return list(
+        tqdm(
+            parallel_map(
+                func=func,
+                input_values=input_values,
+                chunk_size=chunk_size,
+                concurrency=concurrency,
+            ),
+            total=len(input_values),
+            dynamic_ncols=True,
+        )
+    )