From f13f5210bced1ba5915fa1b412417c541471873c Mon Sep 17 00:00:00 2001
From: Benjamin Piwowarski <benjamin@piwowarski.fr>
Date: Wed, 28 Feb 2024 15:55:30 +0100
Subject: [PATCH] Re-arrange utils

---
 src/xpmir/utils/iter.py    |  9 +++++++--
 src/xpmir/utils/logging.py | 29 +++++++++++++++++++++++++++++
 src/xpmir/utils/utils.py   | 29 +----------------------------
 3 files changed, 37 insertions(+), 30 deletions(-)
 create mode 100644 src/xpmir/utils/logging.py

diff --git a/src/xpmir/utils/iter.py b/src/xpmir/utils/iter.py
index bf9f759..597e441 100644
--- a/src/xpmir/utils/iter.py
+++ b/src/xpmir/utils/iter.py
@@ -333,8 +333,8 @@ def __init__(self, iterator: Iterator[T], maxsize=100):
         self.maxsize = maxsize
         self.iterator = iterator
 
-    def __next__(self):
-        # (1) Start a process if needed
+    def start(self):
+        """Start the iterator process"""
         if self.process is None:
             self.queue = mp.Queue(self.maxsize)
             self.process = mp.Process(
@@ -346,6 +346,11 @@ def __next__(self):
             # Start, and register a kill switch
             self.process.start()
             atexit.register(self.kill_subprocess)
+        return self
+
+    def __next__(self):
+        # Start a process if needed
+        self.start()
 
         # Get the next element
         element = self.queue.get()
diff --git a/src/xpmir/utils/logging.py b/src/xpmir/utils/logging.py
new file mode 100644
index 0000000..4e5728a
--- /dev/null
+++ b/src/xpmir/utils/logging.py
@@ -0,0 +1,29 @@
+import logging
+from functools import cached_property
+
+
+def easylog():
+    """
+    Returns a logger with the caller's __name__
+    """
+    import inspect
+
+    try:
+        frame = inspect.stack()[1]  # caller
+        module = inspect.getmodule(frame[0])
+        return logging.getLogger(module.__name__)
+    except IndexError:
+        return logging.getLogger("UNKNOWN")
+
+
+class EasyLogger:
+    @cached_property
+    def logger(self):
+        clsdict = self.__class__.__dict__
+
+        logger = clsdict.get("__LOGGER__", None)
+        if logger is None:
+            logger = logging.getLogger(self.__class__.__qualname__)
+            self.__class__.__LOGGER__ = logger
+
+        return logger
diff --git a/src/xpmir/utils/utils.py b/src/xpmir/utils/utils.py
index dc68ef9..36278a2 100644
--- a/src/xpmir/utils/utils.py
+++ b/src/xpmir/utils/utils.py
@@ -12,7 +12,6 @@
 )
 import inspect
 import logging
-from functools import cached_property
 import os
 from pathlib import Path
 import re
@@ -21,6 +20,7 @@
 from experimaestro import SubmitHook, Job, Launcher
 from threading import Thread
 from xpmir.utils.functools import cache
+from xpmir.utils.logging import easylog, EasyLogger  # noqa: F401
 
 T = TypeVar("T")
 
@@ -143,33 +143,6 @@ def batchiter(batchsize: int, iter: Iterator[T], keeppartial=True) -> Iterator[L
         yield samples
 
 
-def easylog():
-    """
-    Returns a logger with the caller's __name__
-    """
-    import inspect
-
-    try:
-        frame = inspect.stack()[1]  # caller
-        module = inspect.getmodule(frame[0])
-        return logging.getLogger(module.__name__)
-    except IndexError:
-        return logging.getLogger("UNKNOWN")
-
-
-class EasyLogger:
-    @cached_property
-    def logger(self):
-        clsdict = self.__class__.__dict__
-
-        logger = clsdict.get("__LOGGER__", None)
-        if logger is None:
-            logger = logging.getLogger(self.__class__.__qualname__)
-            self.__class__.__LOGGER__ = logger
-
-        return logger
-
-
 @cache
 def find_java_home(min_version: int = 6) -> str:
     """Find JAVA HOME"""