From a115539931b4ca03f4009fcfaccc1a4c18517e58 Mon Sep 17 00:00:00 2001
From: J08nY
Date: Wed, 3 Apr 2024 14:37:17 +0200
Subject: Add utility task executor to make notebooks more readable.

---
 pyecsca/misc/utils.py | 25 ++++++++++++++++++++++++-
 1 file changed, 24 insertions(+), 1 deletion(-)

(limited to 'pyecsca/misc')

diff --git a/pyecsca/misc/utils.py b/pyecsca/misc/utils.py
index 4e6a26b..aa450e3 100644
--- a/pyecsca/misc/utils.py
+++ b/pyecsca/misc/utils.py
@@ -1,9 +1,12 @@
 """Just some utilities I promise."""
 import sys
 from ast import parse
+from typing import List, Any, Generator
 
 from ..misc.cfg import getconfig
 
+from concurrent.futures import ProcessPoolExecutor, as_completed, Future
+
 
 def pexec(s):
     return parse(s, mode="exec")
@@ -17,7 +20,8 @@ def in_notebook() -> bool:
     """Test whether we are executing in Jupyter notebook."""
     try:
         from IPython import get_ipython
-        if 'IPKernelApp' not in get_ipython().config:  # pragma: no cover
+
+        if "IPKernelApp" not in get_ipython().config:  # pragma: no cover
             return False
     except ImportError:
         return False
@@ -36,3 +40,22 @@ def warn(*args, **kwargs):
     """Log a message."""
     if in_notebook() and getconfig().log.enabled:
         print(*args, **kwargs, file=sys.stderr)
+
+
+class TaskExecutor(ProcessPoolExecutor):
+    """A simple ProcessPoolExecutor that keeps tracks of tasks that were submitted to it."""
+    keys: List[Any]
+    futures: List[Future]
+
+    def submit_task(self, key: Any, fn, /, *args, **kwargs):
+        self.futures.append(self.submit(fn, *args, **kwargs))
+        self.keys.append(key)
+
+    @property
+    def tasks(self):
+        return list(zip(self.keys, self.futures))
+
+    def as_completed(self) -> Generator[tuple[Any, Future], Any, None]:
+        for future in as_completed(self.futures):
+            i = self.futures.index(future)
+            yield self.keys[i], future
-- 
cgit v1.2.3-70-g09d2