stouputils 1.9.1__tar.gz → 1.9.3__tar.gz

{stouputils-1.9.1 → stouputils-1.9.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: stouputils
-Version: 1.9.1
+Version: 1.9.3
 Summary: Stouputils is a collection of utility modules designed to simplify and enhance the development process. It includes a range of tools for tasks such as execution of doctests, display utilities, decorators, as well as context managers, and many more.
 Project-URL: Homepage, https://github.com/Stoupy51/stouputils
 Project-URL: Issues, https://github.com/Stoupy51/stouputils/issues
@@ -149,7 +149,7 @@ stouputils all_<TAB>    # Completes to: all_doctests
 ├── <a href="https://stoupy51.github.io/stouputils/latest/modules/stouputils.backup.html">backup.py</a>                <span class="comment"># 💾 Utilities for backup management (delta backup, consolidate)</span>
 ├── <a href="https://stoupy51.github.io/stouputils/latest/modules/stouputils.collections.html">collections.py</a>           <span class="comment"># 🧰 Utilities for collection manipulation (unique_list, sort_dict_keys, upsert_in_dataframe, array_to_disk)</span>
 ├── <a href="https://stoupy51.github.io/stouputils/latest/modules/stouputils.ctx.html">ctx.py</a>                   <span class="comment"># 🔇 Context managers (Muffle, LogToFile, MeasureTime, DoNothing)</span>
-├── <a href="https://stoupy51.github.io/stouputils/latest/modules/stouputils.decorators.html">decorators.py</a>            <span class="comment"># 🎯 Decorators (measure_time, handle_error, simple_cache, retry, abstract, deprecated, silent)</span>
+├── <a href="https://stoupy51.github.io/stouputils/latest/modules/stouputils.decorators.html">decorators.py</a>            <span class="comment"># 🎯 Decorators (measure_time, handle_error, timeout, retry, simple_cache, abstract, deprecated, silent)</span>
 ├── <a href="https://stoupy51.github.io/stouputils/latest/modules/stouputils.image.html">image.py</a>                 <span class="comment"># 🖼️ Little utilities for image processing (image_resize, auto_crop, numpy_to_gif, numpy_to_obj)</span>
 ├── <a href="https://stoupy51.github.io/stouputils/latest/modules/stouputils.io.html">io.py</a>                    <span class="comment"># 💾 Utilities for file management (super_json, super_csv, super_copy, super_open, clean_path)</span>
 ├── <a href="https://stoupy51.github.io/stouputils/latest/modules/stouputils.parallel.html">parallel.py</a>              <span class="comment"># 🔀 Utility functions for parallel processing (multiprocessing, multithreading)</span>

{stouputils-1.9.1 → stouputils-1.9.3}/README.md

@@ -107,7 +107,7 @@ stouputils all_<TAB>    # Completes to: all_doctests
 ├── <a href="https://stoupy51.github.io/stouputils/latest/modules/stouputils.backup.html">backup.py</a>                <span class="comment"># 💾 Utilities for backup management (delta backup, consolidate)</span>
 ├── <a href="https://stoupy51.github.io/stouputils/latest/modules/stouputils.collections.html">collections.py</a>           <span class="comment"># 🧰 Utilities for collection manipulation (unique_list, sort_dict_keys, upsert_in_dataframe, array_to_disk)</span>
 ├── <a href="https://stoupy51.github.io/stouputils/latest/modules/stouputils.ctx.html">ctx.py</a>                   <span class="comment"># 🔇 Context managers (Muffle, LogToFile, MeasureTime, DoNothing)</span>
-├── <a href="https://stoupy51.github.io/stouputils/latest/modules/stouputils.decorators.html">decorators.py</a>            <span class="comment"># 🎯 Decorators (measure_time, handle_error, simple_cache, retry, abstract, deprecated, silent)</span>
+├── <a href="https://stoupy51.github.io/stouputils/latest/modules/stouputils.decorators.html">decorators.py</a>            <span class="comment"># 🎯 Decorators (measure_time, handle_error, timeout, retry, simple_cache, abstract, deprecated, silent)</span>
 ├── <a href="https://stoupy51.github.io/stouputils/latest/modules/stouputils.image.html">image.py</a>                 <span class="comment"># 🖼️ Little utilities for image processing (image_resize, auto_crop, numpy_to_gif, numpy_to_obj)</span>
 ├── <a href="https://stoupy51.github.io/stouputils/latest/modules/stouputils.io.html">io.py</a>                    <span class="comment"># 💾 Utilities for file management (super_json, super_csv, super_copy, super_open, clean_path)</span>
 ├── <a href="https://stoupy51.github.io/stouputils/latest/modules/stouputils.parallel.html">parallel.py</a>              <span class="comment"># 🔀 Utility functions for parallel processing (multiprocessing, multithreading)</span>

{stouputils-1.9.1 → stouputils-1.9.3}/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "stouputils"
-version = "1.9.1"
+version = "1.9.3"
 description = "Stouputils is a collection of utility modules designed to simplify and enhance the development process. It includes a range of tools for tasks such as execution of doctests, display utilities, decorators, as well as context managers, and many more."
 readme = "README.md"
 requires-python = ">=3.10"

{stouputils-1.9.1 → stouputils-1.9.3}/stouputils/ctx.py

@@ -5,6 +5,7 @@ This module provides context managers for temporarily silencing output.
 - MeasureTime: Context manager to measure execution time of a code block
 - Muffle: Context manager that temporarily silences output (alternative to stouputils.decorators.silent())
 - DoNothing: Context manager that does nothing (no-op)
+- SetMPStartMethod: Context manager to temporarily set multiprocessing start method
 
 .. image:: https://raw.githubusercontent.com/Stoupy51/stouputils/refs/heads/main/assets/ctx_module.gif
   :alt: stouputils ctx examples
@@ -284,3 +285,51 @@ class DoNothing:
 		""" Exit async context manager (does nothing) """
 		pass
 
+# Context manager to temporarily set multiprocessing start method
+class SetMPStartMethod:
+	""" Context manager to temporarily set multiprocessing start method.
+
+	This context manager allows you to temporarily change the multiprocessing start method
+	and automatically restores the original method when exiting the context.
+
+	Args:
+		start_method (str): The start method to use: "spawn", "fork", or "forkserver"
+
+	Examples:
+		.. code-block:: python
+
+			> import multiprocessing as mp
+			> import stouputils as stp
+			> # Temporarily use spawn method
+			> with stp.SetMPStartMethod("spawn"):
+			> ...     # Your multiprocessing code here
+			> ...     pass
+
+			> # Original method is automatically restored
+	"""
+	def __init__(self, start_method: str | None) -> None:
+		self.start_method: str | None = start_method
+		""" The start method to use """
+		self.old_method: str | None = None
+		""" The original start method to restore """
+
+	def __enter__(self) -> SetMPStartMethod:
+		""" Enter context manager which sets the start method """
+		if self.start_method is None:
+			return self
+		import multiprocessing as mp
+
+		self.old_method = mp.get_start_method(allow_none=True)
+		if self.old_method != self.start_method:
+			mp.set_start_method(self.start_method, force=True)
+		return self
+
+	def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+		""" Exit context manager which restores the original start method """
+		if self.start_method is None:
+			return
+		import multiprocessing as mp
+
+		if self.old_method != self.start_method:
+			mp.set_start_method(self.old_method, force=True)
+

{stouputils-1.9.1 → stouputils-1.9.3}/stouputils/ctx.pyi

@@ -141,3 +141,32 @@ class DoNothing:
         """ Enter async context manager (does nothing) """
     async def __aexit__(self, *excinfo: Any) -> None:
         """ Exit async context manager (does nothing) """
+
+class SetMPStartMethod:
+    ''' Context manager to temporarily set multiprocessing start method.
+
+\tThis context manager allows you to temporarily change the multiprocessing start method
+\tand automatically restores the original method when exiting the context.
+
+\tArgs:
+\t\tstart_method (str): The start method to use: "spawn", "fork", or "forkserver"
+
+\tExamples:
+\t\t.. code-block:: python
+
+\t\t\t> import multiprocessing as mp
+\t\t\t> import stouputils as stp
+\t\t\t> # Temporarily use spawn method
+\t\t\t> with stp.SetMPStartMethod("spawn"):
+\t\t\t> ...     # Your multiprocessing code here
+\t\t\t> ...     pass
+
+\t\t\t> # Original method is automatically restored
+\t'''
+    start_method: str | None
+    old_method: str | None
+    def __init__(self, start_method: str | None) -> None: ...
+    def __enter__(self) -> SetMPStartMethod:
+        """ Enter context manager which sets the start method """
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        """ Exit context manager which restores the original start method """

{stouputils-1.9.1 → stouputils-1.9.3}/stouputils/decorators.py

@@ -3,8 +3,9 @@ This module provides decorators for various purposes:
 
 - measure_time(): Measure the execution time of a function and print it with the given print function
 - handle_error(): Handle an error with different log levels
-- simple_cache(): Easy cache function with parameter caching method
+- timeout(): Raise an exception if the function runs longer than the specified timeout
 - retry(): Retry a function when specific exceptions are raised, with configurable delay and max attempts
+- simple_cache(): Easy cache function with parameter caching method
 - abstract(): Mark a function as abstract, using LogLevels for error handling
 - deprecated(): Mark a function as deprecated, using LogLevels for warning handling
 - silent(): Make a function silent (disable stdout, and stderr if specified) (alternative to stouputils.ctx.Muffle)
@@ -167,67 +168,105 @@ def handle_error(
 		return decorator
 	return decorator(func)
 
-# Easy cache function with parameter caching method
-def simple_cache(
+# Decorator that raises an exception if the function runs too long
+def timeout(
 	func: Callable[..., Any] | None = None,
 	*,
-	method: Literal["str", "pickle"] = "str"
+	seconds: float = 60.0,
+	message: str = ""
 ) -> Callable[..., Any]:
-	""" Decorator that caches the result of a function based on its arguments.
+	""" Decorator that raises a TimeoutError if the function runs longer than the specified timeout.
 
-	The str method is often faster than the pickle method (by a little) but not as accurate with complex objects.
+	Note: This decorator uses SIGALRM on Unix systems, which only works in the main thread.
+	On Windows or in non-main threads, it will fall back to a polling-based approach.
 
 	Args:
-		func   (Callable[..., Any] | None): Function to cache
-		method (Literal["str", "pickle"]):  The method to use for caching.
+		func		(Callable[..., Any] | None):	Function to apply timeout to
+		seconds		(float):						Timeout duration in seconds (default: 60.0)
+		message		(str):							Custom timeout message (default: "Function '{func_name}' timed out after {seconds} seconds")
+
 	Returns:
-		Callable[..., Any]: A decorator that caches the result of a function.
+		Callable[..., Any]: Decorator that enforces timeout on the function
+
+	Raises:
+		TimeoutError: If the function execution exceeds the timeout duration
+
 	Examples:
-		>>> @simple_cache
-		... def test1(a: int, b: int) -> int:
-		...     return a + b
+		>>> @timeout(seconds=2.0)
+		... def slow_function():
+		...     time.sleep(5)
+		>>> slow_function()  # Raises TimeoutError after 2 seconds
+		Traceback (most recent call last):
+			...
+		TimeoutError: Function 'slow_function' timed out after 2.0 seconds
 
-		>>> @simple_cache(method="str")
-		... def test2(a: int, b: int) -> int:
-		...     return a + b
-		>>> test2(1, 2)
-		3
-		>>> test2(1, 2)
-		3
-		>>> test2(3, 4)
-		7
+		>>> @timeout(seconds=1.0, message="Custom timeout message")
+		... def another_slow_function():
+		...     time.sleep(3)
+		>>> another_slow_function()  # Raises TimeoutError after 1 second
+		Traceback (most recent call last):
+			...
+		TimeoutError: Custom timeout message
 	"""
 	def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
-		# Create the cache dict
-		cache_dict: dict[bytes, Any] = {}
-
-		# Create the wrapper
 		@wraps(func)
 		def wrapper(*args: tuple[Any, ...], **kwargs: dict[str, Any]) -> Any:
+			# Build timeout message
+			msg: str = message if message else f"Function '{_get_func_name(func)}' timed out after {seconds} seconds"
 
-			# Get the hashed key
-			if method == "str":
-				hashed: bytes = str(args).encode() + str(kwargs).encode()
-			elif method == "pickle":
-				hashed: bytes = pickle_dumps((args, kwargs))
-			else:
-				raise ValueError("Invalid caching method. Supported methods are 'str' and 'pickle'.")
+			# Try to use signal-based timeout (Unix only, main thread only)
+			try:
+				import signal
+				def timeout_handler(signum: int, frame: Any) -> None:
+					raise TimeoutError(msg)
 
-			# If the key is in the cache, return it
-			if hashed in cache_dict:
-				return cache_dict[hashed]
+				# Set the signal handler and alarm
+				old_handler = signal.signal(signal.SIGALRM, timeout_handler) # type: ignore
+				signal.setitimer(signal.ITIMER_REAL, seconds) # type: ignore
+
+				try:
+					result = func(*args, **kwargs)
+				finally:
+					# Cancel the alarm and restore the old handler
+					signal.setitimer(signal.ITIMER_REAL, 0) # type: ignore
+					signal.signal(signal.SIGALRM, old_handler) # type: ignore
 
-			# Else, call the function and add the result to the cache
-			else:
-				result: Any = func(*args, **kwargs)
-				cache_dict[hashed] = result
 				return result
 
-		# Return the wrapper
-		wrapper.__name__ = _get_wrapper_name("stouputils.decorators.simple_cache", func)
+			except (ValueError, AttributeError) as e:
+				# SIGALRM not available (Windows) or not in main thread
+				# Fall back to polling-based timeout (less precise but portable)
+				import threading
+
+				result_container: list[Any] = []
+				exception_container: list[BaseException] = []
+
+				def target() -> None:
+					try:
+						result_container.append(func(*args, **kwargs))
+					except BaseException as e_2:
+						exception_container.append(e_2)
+
+				thread = threading.Thread(target=target, daemon=True)
+				thread.start()
+				thread.join(timeout=seconds)
+
+				if thread.is_alive():
+					# Thread is still running, timeout occurred
+					raise TimeoutError(msg) from e
+
+				# Check if an exception was raised in the thread
+				if exception_container:
+					raise exception_container[0] from e
+
+				# Return the result if available
+				if result_container:
+					return result_container[0]
+
+		wrapper.__name__ = _get_wrapper_name("stouputils.decorators.timeout", func)
 		return wrapper
 
-	# Handle both @simple_cache and @simple_cache(method=...)
+	# Handle both @timeout and @timeout(seconds=..., message=...)
 	if func is None:
 		return decorator
 	return decorator(func)
@@ -307,6 +346,71 @@ def retry(
 		return decorator
 	return decorator(func)
 
+# Easy cache function with parameter caching method
+def simple_cache(
+	func: Callable[..., Any] | None = None,
+	*,
+	method: Literal["str", "pickle"] = "str"
+) -> Callable[..., Any]:
+	""" Decorator that caches the result of a function based on its arguments.
+
+	The str method is often faster than the pickle method (by a little) but not as accurate with complex objects.
+
+	Args:
+		func   (Callable[..., Any] | None): Function to cache
+		method (Literal["str", "pickle"]):  The method to use for caching.
+	Returns:
+		Callable[..., Any]: A decorator that caches the result of a function.
+	Examples:
+		>>> @simple_cache
+		... def test1(a: int, b: int) -> int:
+		...     return a + b
+
+		>>> @simple_cache(method="str")
+		... def test2(a: int, b: int) -> int:
+		...     return a + b
+		>>> test2(1, 2)
+		3
+		>>> test2(1, 2)
+		3
+		>>> test2(3, 4)
+		7
+	"""
+	def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
+		# Create the cache dict
+		cache_dict: dict[bytes, Any] = {}
+
+		# Create the wrapper
+		@wraps(func)
+		def wrapper(*args: tuple[Any, ...], **kwargs: dict[str, Any]) -> Any:
+
+			# Get the hashed key
+			if method == "str":
+				hashed: bytes = str(args).encode() + str(kwargs).encode()
+			elif method == "pickle":
+				hashed: bytes = pickle_dumps((args, kwargs))
+			else:
+				raise ValueError("Invalid caching method. Supported methods are 'str' and 'pickle'.")
+
+			# If the key is in the cache, return it
+			if hashed in cache_dict:
+				return cache_dict[hashed]
+
+			# Else, call the function and add the result to the cache
+			else:
+				result: Any = func(*args, **kwargs)
+				cache_dict[hashed] = result
+				return result
+
+		# Return the wrapper
+		wrapper.__name__ = _get_wrapper_name("stouputils.decorators.simple_cache", func)
+		return wrapper
+
+	# Handle both @simple_cache and @simple_cache(method=...)
+	if func is None:
+		return decorator
+	return decorator(func)
+
 # Decorator that marks a function as abstract
 def abstract(
 	func: Callable[..., Any] | None = None,

{stouputils-1.9.1 → stouputils-1.9.3}/stouputils/decorators.pyi

@@ -64,30 +64,39 @@ def handle_error(func: Callable[..., Any] | None = None, *, exceptions: tuple[ty
 \t\t...     raise ValueError("Let\'s fail")
 \t\t>>> # test()\t# [WARNING HH:MM:SS] Error during test: (ValueError) Let\'s fail
 \t'''
-def simple_cache(func: Callable[..., Any] | None = None, *, method: Literal['str', 'pickle'] = 'str') -> Callable[..., Any]:
-    ''' Decorator that caches the result of a function based on its arguments.
+def timeout(func: Callable[..., Any] | None = None, *, seconds: float = 60.0, message: str = '') -> Callable[..., Any]:
+    ''' Decorator that raises a TimeoutError if the function runs longer than the specified timeout.
 
-\tThe str method is often faster than the pickle method (by a little) but not as accurate with complex objects.
+\tNote: This decorator uses SIGALRM on Unix systems, which only works in the main thread.
+\tOn Windows or in non-main threads, it will fall back to a polling-based approach.
 
 \tArgs:
-\t\tfunc   (Callable[..., Any] | None): Function to cache
-\t\tmethod (Literal["str", "pickle"]):  The method to use for caching.
+\t\tfunc\t\t(Callable[..., Any] | None):\tFunction to apply timeout to
+\t\tseconds\t\t(float):\t\t\t\t\t\tTimeout duration in seconds (default: 60.0)
+\t\tmessage\t\t(str):\t\t\t\t\t\t\tCustom timeout message (default: "Function \'{func_name}\' timed out after {seconds} seconds")
+
 \tReturns:
-\t\tCallable[..., Any]: A decorator that caches the result of a function.
+\t\tCallable[..., Any]: Decorator that enforces timeout on the function
+
+\tRaises:
+\t\tTimeoutError: If the function execution exceeds the timeout duration
+
 \tExamples:
-\t\t>>> @simple_cache
-\t\t... def test1(a: int, b: int) -> int:
-\t\t...     return a + b
+\t\t>>> @timeout(seconds=2.0)
+\t\t... def slow_function():
+\t\t...     time.sleep(5)
+\t\t>>> slow_function()  # Raises TimeoutError after 2 seconds
+\t\tTraceback (most recent call last):
+\t\t\t...
+\t\tTimeoutError: Function \'slow_function\' timed out after 2.0 seconds
 
-\t\t>>> @simple_cache(method="str")
-\t\t... def test2(a: int, b: int) -> int:
-\t\t...     return a + b
-\t\t>>> test2(1, 2)
-\t\t3
-\t\t>>> test2(1, 2)
-\t\t3
-\t\t>>> test2(3, 4)
-\t\t7
+\t\t>>> @timeout(seconds=1.0, message="Custom timeout message")
+\t\t... def another_slow_function():
+\t\t...     time.sleep(3)
+\t\t>>> another_slow_function()  # Raises TimeoutError after 1 second
+\t\tTraceback (most recent call last):
+\t\t\t...
+\t\tTimeoutError: Custom timeout message
 \t'''
 def retry(func: Callable[..., Any] | None = None, *, exceptions: tuple[type[BaseException], ...] | type[BaseException] = ..., max_attempts: int = 10, delay: float = 1.0, backoff: float = 1.0, message: str = '') -> Callable[..., Any]:
     ''' Decorator that retries a function when specific exceptions are raised.
@@ -118,6 +127,31 @@ def retry(func: Callable[..., Any] | None = None, *, exceptions: tuple[type[Base
 \t\t... def might_fail():
 \t\t...     pass
 \t'''
+def simple_cache(func: Callable[..., Any] | None = None, *, method: Literal['str', 'pickle'] = 'str') -> Callable[..., Any]:
+    ''' Decorator that caches the result of a function based on its arguments.
+
+\tThe str method is often faster than the pickle method (by a little) but not as accurate with complex objects.
+
+\tArgs:
+\t\tfunc   (Callable[..., Any] | None): Function to cache
+\t\tmethod (Literal["str", "pickle"]):  The method to use for caching.
+\tReturns:
+\t\tCallable[..., Any]: A decorator that caches the result of a function.
+\tExamples:
+\t\t>>> @simple_cache
+\t\t... def test1(a: int, b: int) -> int:
+\t\t...     return a + b
+
+\t\t>>> @simple_cache(method="str")
+\t\t... def test2(a: int, b: int) -> int:
+\t\t...     return a + b
+\t\t>>> test2(1, 2)
+\t\t3
+\t\t>>> test2(1, 2)
+\t\t3
+\t\t>>> test2(3, 4)
+\t\t7
+\t'''
 def abstract(func: Callable[..., Any] | None = None, *, error_log: LogLevels = ...) -> Callable[..., Any]:
     """ Decorator that marks a function as abstract.
 

{stouputils-1.9.1 → stouputils-1.9.3}/stouputils/parallel.py

@@ -17,7 +17,7 @@ import time
 from collections.abc import Callable
 from typing import Any, TypeVar, cast
 
-from .decorators import LogLevels, handle_error
+from .ctx import SetMPStartMethod
 from .print import BAR_FORMAT, MAGENTA
 
 
@@ -34,9 +34,8 @@ T = TypeVar("T")
 R = TypeVar("R")
 
 # Functions
-@handle_error(error_log=LogLevels.ERROR_TRACEBACK)
 def multiprocessing(
-	func: Callable[[T], R] | list[Callable[[T], R]],
+	func: Callable[..., R] | list[Callable[..., R]],
 	args: list[T],
 	use_starmap: bool = False,
 	chunksize: int = 1,
@@ -131,14 +130,8 @@ def multiprocessing(
 			if "SemLock created in a fork context is being shared with a process in a spawn context" in str(e):
 
 				# Try with alternate start method
-				old_method: str | None = mp.get_start_method(allow_none=True)
-				new_method: str = "spawn" if old_method in (None, "fork") else "fork"
-				mp.set_start_method(new_method, force=True)
-
-				try:
+				with SetMPStartMethod("spawn" if mp.get_start_method() != "spawn" else "fork"):
 					return process()
-				finally:
-					mp.set_start_method(old_method, force=True)
 			else: # Re-raise if it's not the SemLock error
 				raise
 
@@ -150,9 +143,8 @@ def multiprocessing(
 			return [func(arg) for arg in args]
 
 
-@handle_error(error_log=LogLevels.ERROR_TRACEBACK)
 def multithreading(
-	func: Callable[[T], R] | list[Callable[[T], R]],
+	func: Callable[..., R] | list[Callable[..., R]],
 	args: list[T],
 	use_starmap: bool = False,
 	desc: str = "",
@@ -243,10 +235,10 @@ def multithreading(
 			return [func(arg) for arg in args]
 
 
-@handle_error(error_log=LogLevels.ERROR_TRACEBACK)
 def run_in_subprocess(
 	func: Callable[..., R],
 	*args: Any,
+	timeout: float | None = None,
 	**kwargs: Any
 ) -> R:
 	""" Execute a function in a subprocess with positional and keyword arguments.
@@ -256,16 +248,19 @@ def run_in_subprocess(
 	be created, run the function with the provided arguments, and return the result.
 
 	Args:
-		func     (Callable): The function to execute in a subprocess.
+		func         (Callable):     The function to execute in a subprocess.
 			(SHOULD BE A TOP-LEVEL FUNCTION TO BE PICKLABLE)
-		*args    (Any):      Positional arguments to pass to the function.
-		**kwargs (Any):      Keyword arguments to pass to the function.
+		*args        (Any):          Positional arguments to pass to the function.
+		timeout      (float | None): Maximum time in seconds to wait for the subprocess.
+			If None, wait indefinitely. If the subprocess exceeds this time, it will be terminated.
+		**kwargs     (Any):          Keyword arguments to pass to the function.
 
 	Returns:
 		R: The return value of the function.
 
 	Raises:
-		RuntimeError: If the subprocess exits with a non-zero exit code.
+		RuntimeError: If the subprocess exits with a non-zero exit code or times out.
+		TimeoutError: If the subprocess exceeds the specified timeout.
 
 	Examples:
 		.. code-block:: python
@@ -285,6 +280,9 @@ def run_in_subprocess(
 			.     return f"{greeting}, {name}!"
 			> run_in_subprocess(greet, "World", greeting="Hi")
 			'Hi, World!'
+
+			> # With timeout to prevent hanging
+			> run_in_subprocess(some_gpu_func, data, timeout=300.0)
 	"""
 	import multiprocessing as mp
 	from multiprocessing import Queue
@@ -298,20 +296,40 @@ def run_in_subprocess(
 		args=(result_queue, func, args, kwargs)
 	)
 	process.start()
-	process.join()
+
+	# Join with timeout to prevent indefinite hanging
+	process.join(timeout=timeout)
+
+	# Check if process is still alive (timed out)
+	if process.is_alive():
+		process.terminate()
+		time.sleep(0.5)  # Give it a moment to terminate gracefully
+		if process.is_alive():
+			process.kill()
+		process.join()
+		raise TimeoutError(f"Subprocess exceeded timeout of {timeout} seconds and was terminated")
 
 	# Check exit code
 	if process.exitcode != 0:
-		raise RuntimeError(f"Subprocess failed with exit code {process.exitcode}")
+		# Try to get any exception from the queue (non-blocking)
+		error_msg = f"Subprocess failed with exit code {process.exitcode}"
+		try:
+			if not result_queue.empty():
+				result_or_exception = result_queue.get_nowait()
+				if isinstance(result_or_exception, Exception):
+					raise result_or_exception
+		except Exception:
+			pass
+		raise RuntimeError(error_msg)
 
 	# Retrieve the result
-	if not result_queue.empty():
-		result_or_exception = result_queue.get()
+	try:
+		result_or_exception = result_queue.get_nowait()
 		if isinstance(result_or_exception, Exception):
 			raise result_or_exception
 		return result_or_exception
-	else:
-		raise RuntimeError("Subprocess did not return any result")
+	except Exception as e:
+		raise RuntimeError("Subprocess did not return any result") from e
 
 
 # "Private" function for subprocess wrapper (must be at module level for pickling on Windows)

{stouputils-1.9.1 → stouputils-1.9.3}/stouputils/parallel.pyi

@@ -1,4 +1,4 @@
-from .decorators import LogLevels as LogLevels, handle_error as handle_error
+from .ctx import SetMPStartMethod as SetMPStartMethod
 from .print import BAR_FORMAT as BAR_FORMAT, MAGENTA as MAGENTA
 from collections.abc import Callable
 from typing import Any, TypeVar
@@ -10,7 +10,7 @@ CPU_COUNT: int
 T = TypeVar('T')
 R = TypeVar('R')
 
-def multiprocessing(func: Callable[[T], R] | list[Callable[[T], R]], args: list[T], use_starmap: bool = False, chunksize: int = 1, desc: str = '', max_workers: int = ..., delay_first_calls: float = 0, color: str = ..., bar_format: str = ..., ascii: bool = False) -> list[R]:
+def multiprocessing(func: Callable[..., R] | list[Callable[..., R]], args: list[T], use_starmap: bool = False, chunksize: int = 1, desc: str = '', max_workers: int = ..., delay_first_calls: float = 0, color: str = ..., bar_format: str = ..., ascii: bool = False) -> list[R]:
     ''' Method to execute a function in parallel using multiprocessing
 
 \t- For CPU-bound operations where the GIL (Global Interpreter Lock) is a bottleneck.
@@ -64,7 +64,7 @@ def multiprocessing(func: Callable[[T], R] | list[Callable[[T], R]], args: list[
 \t\t\t. )
 \t\t\t[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
 \t'''
-def multithreading(func: Callable[[T], R] | list[Callable[[T], R]], args: list[T], use_starmap: bool = False, desc: str = '', max_workers: int = ..., delay_first_calls: float = 0, color: str = ..., bar_format: str = ..., ascii: bool = False) -> list[R]:
+def multithreading(func: Callable[..., R] | list[Callable[..., R]], args: list[T], use_starmap: bool = False, desc: str = '', max_workers: int = ..., delay_first_calls: float = 0, color: str = ..., bar_format: str = ..., ascii: bool = False) -> list[R]:
     ''' Method to execute a function in parallel using multithreading, you should use it:
 
 \t- For I/O-bound operations where the GIL is not a bottleneck, such as network requests or disk operations.
@@ -116,7 +116,7 @@ def multithreading(func: Callable[[T], R] | list[Callable[[T], R]], args: list[T
 \t\t\t. )
 \t\t\t[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
 \t'''
-def run_in_subprocess(func: Callable[..., R], *args: Any, **kwargs: Any) -> R:
+def run_in_subprocess(func: Callable[..., R], *args: Any, timeout: float | None = None, **kwargs: Any) -> R:
     ''' Execute a function in a subprocess with positional and keyword arguments.
 
 \tThis is useful when you need to run a function in isolation to avoid memory leaks,
@@ -124,16 +124,19 @@ def run_in_subprocess(func: Callable[..., R], *args: Any, **kwargs: Any) -> R:
 \tbe created, run the function with the provided arguments, and return the result.
 
 \tArgs:
-\t\tfunc     (Callable): The function to execute in a subprocess.
+\t\tfunc         (Callable):     The function to execute in a subprocess.
 \t\t\t(SHOULD BE A TOP-LEVEL FUNCTION TO BE PICKLABLE)
-\t\t*args    (Any):      Positional arguments to pass to the function.
-\t\t**kwargs (Any):      Keyword arguments to pass to the function.
+\t\t*args        (Any):          Positional arguments to pass to the function.
+\t\ttimeout      (float | None): Maximum time in seconds to wait for the subprocess.
+\t\t\tIf None, wait indefinitely. If the subprocess exceeds this time, it will be terminated.
+\t\t**kwargs     (Any):          Keyword arguments to pass to the function.
 
 \tReturns:
 \t\tR: The return value of the function.
 
 \tRaises:
-\t\tRuntimeError: If the subprocess exits with a non-zero exit code.
+\t\tRuntimeError: If the subprocess exits with a non-zero exit code or times out.
+\t\tTimeoutError: If the subprocess exceeds the specified timeout.
 
 \tExamples:
 \t\t.. code-block:: python
@@ -153,6 +156,9 @@ def run_in_subprocess(func: Callable[..., R], *args: Any, **kwargs: Any) -> R:
 \t\t\t.     return f"{greeting}, {name}!"
 \t\t\t> run_in_subprocess(greet, "World", greeting="Hi")
 \t\t\t\'Hi, World!\'
+
+\t\t\t> # With timeout to prevent hanging
+\t\t\t> run_in_subprocess(some_gpu_func, data, timeout=300.0)
 \t'''
 def _subprocess_wrapper(result_queue: Any, func: Callable[..., R], args: tuple[Any, ...], kwargs: dict[str, Any]) -> None:
     """ Wrapper function to execute the target function and store the result in the queue.