stouputils 1.16.3__py3-none-any.whl → 1.18.0__py3-none-any.whl

stouputils/parallel/__init__.py

@@ -0,0 +1,29 @@
+"""
+This module provides utility functions for parallel processing, such as:
+
+- multiprocessing(): Execute a function in parallel using multiprocessing
+- multithreading(): Execute a function in parallel using multithreading
+- run_in_subprocess(): Execute a function in a subprocess with args and kwargs
+
+I highly encourage you to read the function docstrings to understand when to use each method.
+
+Priority (nice) mapping for multiprocessing():
+
+- Unix-style values from -20 (highest priority) to 19 (lowest priority)
+- Windows automatic mapping:
+  * -20 to -10: HIGH_PRIORITY_CLASS
+  * -9 to -1: ABOVE_NORMAL_PRIORITY_CLASS
+  * 0: NORMAL_PRIORITY_CLASS
+  * 1 to 9: BELOW_NORMAL_PRIORITY_CLASS
+  * 10 to 19: IDLE_PRIORITY_CLASS
+
+.. image:: https://raw.githubusercontent.com/Stoupy51/stouputils/refs/heads/main/assets/parallel_module.gif
+  :alt: stouputils parallel examples
+"""
+
+# Imports
+from .capturer import *
+from .common import *
+from .multi import *
+from .subprocess import *
+

stouputils/parallel/__init__.pyi

@@ -0,0 +1,4 @@
+from .capturer import *
+from .common import *
+from .multi import *
+from .subprocess import *

stouputils/parallel/capturer.py

@@ -0,0 +1,133 @@
+
+# Imports
+import os
+from typing import IO, Any
+
+from ..io import safe_close
+
+
+class PipeWriter:
+	""" A writer that sends data to a multiprocessing Connection. """
+	def __init__(self, conn: Any, encoding: str, errors: str):
+		self.conn: Any = conn
+		self.encoding: str = encoding
+		self.errors: str = errors
+
+	def write(self, data: str) -> int:
+		self.conn.send_bytes(data.encode(self.encoding, errors=self.errors))
+		return len(data)
+
+	def flush(self) -> None:
+		pass
+
+
+class CaptureOutput:
+	""" Utility to capture stdout/stderr from a subprocess and relay it to the parent's stdout.
+
+	The class creates an os.pipe(), marks fds as inheritable (for spawn method),
+	provides methods to start a listener thread that reads from the pipe and writes
+	to the main process's sys.stdout/sys.stderr, and to close/join the listener.
+	"""
+	def __init__(self, encoding: str = "utf-8", errors: str = "replace", chunk_size: int = 1024):
+		import multiprocessing as mp
+		import threading
+		self.encoding: str = encoding
+		self.errors: str = errors
+		self.chunk_size: int = chunk_size
+		self.read_conn, self.write_conn = mp.Pipe(duplex=False)
+		self.read_fd = self.read_conn.fileno()
+		self.write_fd = self.write_conn.fileno()
+		# Internal state for the listener thread and reader handle
+		self._thread: threading.Thread | None = None
+		self._reader_file: IO[Any] | None = None
+		# Sentinel string that will terminate the listener when seen in the stream
+		try:
+			os.set_inheritable(self.read_fd, True)
+			os.set_inheritable(self.write_fd, True)
+		except Exception:
+			pass
+
+	def __repr__(self) -> str:
+		return f"<CaptureOutput read_fd={self.read_fd} write_fd={self.write_fd}>"
+
+	# Pickle support: exclude unpicklable attributes
+	def __getstate__(self) -> dict[str, Any]:
+		state = self.__dict__.copy()
+		state["_thread"] = None
+		return state
+
+	def redirect(self) -> None:
+		""" Redirect sys.stdout and sys.stderr to the pipe's write end. """
+		import sys
+		writer = PipeWriter(self.write_conn, self.encoding, self.errors)
+		sys.stdout = writer
+		sys.stderr = writer
+
+	def parent_close_write(self) -> None:
+		""" Close the parent's copy of the write end; the child's copy remains. """
+		safe_close(self.write_fd)
+		self.write_conn.close()
+		self.write_fd = -1	# Prevent accidental reuse
+
+	def start_listener(self) -> None:
+		""" Start a daemon thread that forwards data from the pipe to sys.stdout/sys.stderr. """
+		import sys
+		if self._thread is not None:
+			return
+
+		# Handler function for reading from the pipe
+		buffer: str = ""
+		def _handle_buffer() -> None:
+			nonlocal buffer
+			if buffer:
+				try:
+					sys.stdout.write(buffer)
+					sys.stdout.flush()
+				except Exception:
+					pass
+				buffer = ""
+
+		# Thread target function
+		def _reader() -> None:
+			nonlocal buffer
+			try:
+				while True:
+					# Read a chunk from the pipe, stop loop on error
+					try:
+						data: bytes = self.read_conn.recv_bytes(self.chunk_size)
+					except EOFError:
+						_handle_buffer()
+						break
+
+					# Decode bytes to text & append to buffer
+					try:
+						chunk: str = data.decode(self.encoding, errors=self.errors)
+					except Exception:
+						chunk = data.decode(self.encoding, errors="replace")
+					buffer += chunk
+
+					# Periodically flush large buffers to avoid holding too much memory
+					if len(buffer) > self.chunk_size * 4:
+						_handle_buffer()
+			finally:
+				safe_close(self.read_fd)
+				self.read_conn.close()
+				self.read_fd = -1
+				self._thread = None		# Mark thread as stopped so callers don't block unnecessarily
+
+		# Start the listener thread
+		import threading
+		self._thread = threading.Thread(target=_reader, daemon=True)
+		self._thread.start()
+
+	def join_listener(self, timeout: float | None = None) -> None:
+		""" Wait for the listener thread to finish (until EOF). """
+		if self._thread is None:
+			safe_close(self.read_fd)
+			return self.read_conn.close()
+		self._thread.join(timeout)
+
+		# If thread finished, ensure read fd is closed and clear thread
+		if self._thread and not self._thread.is_alive():
+			self._thread = None
+

stouputils/parallel/capturer.pyi

@@ -0,0 +1,38 @@
+from ..io import safe_close as safe_close
+from _typeshed import Incomplete
+from typing import Any, IO
+
+class PipeWriter:
+    """ A writer that sends data to a multiprocessing Connection. """
+    conn: Any
+    encoding: str
+    errors: str
+    def __init__(self, conn: Any, encoding: str, errors: str) -> None: ...
+    def write(self, data: str) -> int: ...
+    def flush(self) -> None: ...
+
+class CaptureOutput:
+    """ Utility to capture stdout/stderr from a subprocess and relay it to the parent's stdout.
+
+\tThe class creates an os.pipe(), marks fds as inheritable (for spawn method),
+\tprovides methods to start a listener thread that reads from the pipe and writes
+\tto the main process's sys.stdout/sys.stderr, and to close/join the listener.
+\t"""
+    encoding: str
+    errors: str
+    chunk_size: int
+    read_fd: Incomplete
+    write_fd: Incomplete
+    _thread: threading.Thread | None
+    _reader_file: IO[Any] | None
+    def __init__(self, encoding: str = 'utf-8', errors: str = 'replace', chunk_size: int = 1024) -> None: ...
+    def __repr__(self) -> str: ...
+    def __getstate__(self) -> dict[str, Any]: ...
+    def redirect(self) -> None:
+        """ Redirect sys.stdout and sys.stderr to the pipe's write end. """
+    def parent_close_write(self) -> None:
+        """ Close the parent's copy of the write end; the child's copy remains. """
+    def start_listener(self) -> None:
+        """ Start a daemon thread that forwards data from the pipe to sys.stdout/sys.stderr. """
+    def join_listener(self, timeout: float | None = None) -> None:
+        """ Wait for the listener thread to finish (until EOF). """

stouputils/parallel/common.py

@@ -0,0 +1,134 @@
+
+# Imports
+import os
+import time
+from collections.abc import Callable
+from typing import cast
+
+# Constants
+CPU_COUNT: int = cast(int, os.cpu_count())
+
+
+# "Private" function to wrap function execution with nice priority (must be at module level for pickling)
+def nice_wrapper[T, R](args: tuple[int, Callable[[T], R], T]) -> R:
+	""" Wrapper that applies nice priority then executes the function.
+
+	Args:
+		args (tuple): Tuple containing (nice_value, func, arg)
+
+	Returns:
+		R: Result of the function execution
+	"""
+	nice_value, func, arg = args
+	set_process_priority(nice_value)
+	return func(arg)
+
+# "Private" function to set process priority (must be at module level for pickling on Windows)
+def set_process_priority(nice_value: int) -> None:
+	""" Set the priority of the current process.
+
+	Args:
+		nice_value (int): Unix-style priority value (-20 to 19)
+	"""
+	try:
+		import sys
+		if sys.platform == "win32":
+			# Map Unix nice values to Windows priority classes
+			# -20 to -10: HIGH, -9 to -1: ABOVE_NORMAL, 0: NORMAL, 1-9: BELOW_NORMAL, 10-19: IDLE
+			import ctypes
+			# Windows priority class constants
+			if nice_value <= -10:
+				priority = 0x00000080  # HIGH_PRIORITY_CLASS
+			elif nice_value < 0:
+				priority = 0x00008000  # ABOVE_NORMAL_PRIORITY_CLASS
+			elif nice_value == 0:
+				priority = 0x00000020  # NORMAL_PRIORITY_CLASS
+			elif nice_value < 10:
+				priority = 0x00004000  # BELOW_NORMAL_PRIORITY_CLASS
+			else:
+				priority = 0x00000040  # IDLE_PRIORITY_CLASS
+			kernel32 = ctypes.windll.kernel32
+			handle = kernel32.GetCurrentProcess()
+			kernel32.SetPriorityClass(handle, priority)
+		else:
+			# Unix-like systems
+			os.nice(nice_value)
+	except Exception:
+		pass  # Silently ignore if we can't set priority
+
+# "Private" function to use starmap using args[0](*args[1])
+def starmap[T, R](args: tuple[Callable[[T], R], list[T]]) -> R:
+	r""" Private function to use starmap using args[0](\*args[1])
+
+	Args:
+		args (tuple): Tuple containing the function and the arguments list to pass to the function
+	Returns:
+		object: Result of the function execution
+	"""
+	func, arguments = args
+	return func(*arguments)
+
+# "Private" function to apply delay before calling the target function
+def delayed_call[T, R](args: tuple[Callable[[T], R], float, T]) -> R:
+	""" Private function to apply delay before calling the target function
+
+	Args:
+		args (tuple): Tuple containing the function, delay in seconds, and the argument to pass to the function
+	Returns:
+		object: Result of the function execution
+	"""
+	func, delay, arg = args
+	time.sleep(delay)
+	return func(arg)
+
+# "Private" function to handle parameters for multiprocessing or multithreading functions
+def handle_parameters[T, R](
+	func: Callable[[T], R] | list[Callable[[T], R]],
+	args: list[T],
+	use_starmap: bool,
+	delay_first_calls: float,
+	max_workers: int,
+	desc: str,
+	color: str
+) -> tuple[str, Callable[[T], R], list[T]]:
+	r""" Private function to handle the parameters for multiprocessing or multithreading functions
+
+	Args:
+		func				(Callable | list[Callable]):	Function to execute, or list of functions (one per argument)
+		args				(list):				List of arguments to pass to the function(s)
+		use_starmap			(bool):				Whether to use starmap or not (Defaults to False):
+			True means the function will be called like func(\*args[i]) instead of func(args[i])
+		delay_first_calls	(int):				Apply i*delay_first_calls seconds delay to the first "max_workers" calls.
+			For instance, the first process will be delayed by 0 seconds, the second by 1 second, etc. (Defaults to 0):
+			This can be useful to avoid functions being called in the same second.
+		max_workers			(int):				Number of workers to use
+		desc				(str):				Description of the function execution displayed in the progress bar
+		color				(str):				Color of the progress bar
+
+	Returns:
+		tuple[str, Callable[[T], R], list[T]]:	Tuple containing the description, function, and arguments
+	"""
+	desc = color + desc
+
+	# Handle list of functions: validate and convert to starmap format
+	if isinstance(func, list):
+		func = cast(list[Callable[[T], R]], func)
+		assert len(func) == len(args), f"Length mismatch: {len(func)} functions but {len(args)} arguments"
+		args = [(f, arg if use_starmap else (arg,)) for f, arg in zip(func, args, strict=False)] # type: ignore
+		func = starmap # type: ignore
+
+	# If use_starmap is True, we use the _starmap function
+	elif use_starmap:
+		args = [(func, arg) for arg in args] # type: ignore
+		func = starmap # type: ignore
+
+	# Prepare delayed function calls if delay_first_calls is set
+	if delay_first_calls > 0:
+		args = [
+			(func, i * delay_first_calls if i < max_workers else 0, arg) # type: ignore
+			for i, arg in enumerate(args)
+		]
+		func = delayed_call  # type: ignore
+
+	return desc, func, args # type: ignore
+

stouputils/parallel/common.pyi

@@ -0,0 +1,53 @@
+from collections.abc import Callable
+
+CPU_COUNT: int
+
+def nice_wrapper[T, R](args: tuple[int, Callable[[T], R], T]) -> R:
+    """ Wrapper that applies nice priority then executes the function.
+
+\tArgs:
+\t\targs (tuple): Tuple containing (nice_value, func, arg)
+
+\tReturns:
+\t\tR: Result of the function execution
+\t"""
+def set_process_priority(nice_value: int) -> None:
+    """ Set the priority of the current process.
+
+\tArgs:
+\t\tnice_value (int): Unix-style priority value (-20 to 19)
+\t"""
+def starmap[T, R](args: tuple[Callable[[T], R], list[T]]) -> R:
+    """ Private function to use starmap using args[0](\\*args[1])
+
+\tArgs:
+\t\targs (tuple): Tuple containing the function and the arguments list to pass to the function
+\tReturns:
+\t\tobject: Result of the function execution
+\t"""
+def delayed_call[T, R](args: tuple[Callable[[T], R], float, T]) -> R:
+    """ Private function to apply delay before calling the target function
+
+\tArgs:
+\t\targs (tuple): Tuple containing the function, delay in seconds, and the argument to pass to the function
+\tReturns:
+\t\tobject: Result of the function execution
+\t"""
+def handle_parameters[T, R](func: Callable[[T], R] | list[Callable[[T], R]], args: list[T], use_starmap: bool, delay_first_calls: float, max_workers: int, desc: str, color: str) -> tuple[str, Callable[[T], R], list[T]]:
+    ''' Private function to handle the parameters for multiprocessing or multithreading functions
+
+\tArgs:
+\t\tfunc\t\t\t\t(Callable | list[Callable]):\tFunction to execute, or list of functions (one per argument)
+\t\targs\t\t\t\t(list):\t\t\t\tList of arguments to pass to the function(s)
+\t\tuse_starmap\t\t\t(bool):\t\t\t\tWhether to use starmap or not (Defaults to False):
+\t\t\tTrue means the function will be called like func(\\*args[i]) instead of func(args[i])
+\t\tdelay_first_calls\t(int):\t\t\t\tApply i*delay_first_calls seconds delay to the first "max_workers" calls.
+\t\t\tFor instance, the first process will be delayed by 0 seconds, the second by 1 second, etc. (Defaults to 0):
+\t\t\tThis can be useful to avoid functions being called in the same second.
+\t\tmax_workers\t\t\t(int):\t\t\t\tNumber of workers to use
+\t\tdesc\t\t\t\t(str):\t\t\t\tDescription of the function execution displayed in the progress bar
+\t\tcolor\t\t\t\t(str):\t\t\t\tColor of the progress bar
+
+\tReturns:
+\t\ttuple[str, Callable[[T], R], list[T]]:\tTuple containing the description, function, and arguments
+\t'''