stouputils 1.14.3__py3-none-any.whl → 1.15.1__py3-none-any.whl

stouputils/installer/windows.py

@@ -1,136 +1,136 @@
-""" Installer module for Windows specific functions.
-
-Provides Windows specific implementations for checking administrator privileges,
-determining appropriate installation paths (global/local), and modifying
-the user's PATH environment variable.
-"""
-# Imports
-import os
-
-from ..decorators import LogLevels, handle_error
-from ..io import clean_path
-from ..print import debug, info, warning
-from .common import ask_install_type, prompt_for_path
-
-
-# Functions
-@handle_error(message="Failed to add to PATH (Windows)", error_log=LogLevels.WARNING_TRACEBACK)
-def add_to_path_windows(install_path: str) -> bool | None:
-	""" Add install_path to the User PATH environment variable on Windows.
-
-	Args:
-		install_path (str): The path to add to the User PATH environment variable.
-
-	Returns:
-		bool | None: True if the path was added to the User PATH environment variable, None otherwise.
-	"""
-	# Convert install_path to a Windows path if it's not already
-	install_path = install_path.replace("/", "\\")
-	os.makedirs(install_path, exist_ok=True)
-
-	# Get current user PATH
-	import winreg
-	with winreg.OpenKey(winreg.HKEY_CURRENT_USER, "Environment", 0, winreg.KEY_READ | winreg.KEY_WRITE) as key:
-
-		# Get the number of values in the registry key
-		num_values = winreg.QueryInfoKey(key)[1]
-
-		# Find the index of the 'Path' value
-		path_index = -1
-		for i in range(num_values):
-			if winreg.EnumValue(key, i)[0] == 'Path':
-				path_index = i
-				break
-
-		# Get the current path value
-		current_path: str = winreg.EnumValue(key, path_index)[1]
-
-		# Check if path is already present
-		if install_path not in current_path.split(';'):
-			new_path: str = f"{current_path};{install_path}"
-			winreg.SetValueEx(key, "Path", 0, winreg.REG_EXPAND_SZ, new_path)
-			debug(f"Added '{install_path}' to user PATH. Please restart your terminal for changes to take effect.")
-		else:
-			debug(f"'{install_path}' is already in user PATH.")
-	return True
-
-
-def check_admin_windows() -> bool:
-	""" Check if the script is running with administrator privileges on Windows. """
-	try:
-		import ctypes
-		return ctypes.windll.shell32.IsUserAnAdmin() != 0
-	except Exception:
-		return False
-
-
-@handle_error(message="Failed to get installation path (Windows)", error_log=LogLevels.ERROR_TRACEBACK)
-def get_install_path_windows(
-	program_name: str,
-	ask_global: int = 0,
-	add_path: bool = True,
-	append_to_path: str = "",
-	default_global: str = os.environ.get("ProgramFiles", "C:\\Program Files")
-) -> str:
-	""" Get the installation path for the program
-
-	Args:
-		program_name  (str):   The name of the program to install.
-		ask_global    (int):   0 = ask for anything, 1 = install globally, 2 = install locally
-		add_path      (bool):  Whether to add the program to the PATH environment variable. (Only if installed globally)
-		append_to_path (str):  String to append to the installation path when adding to PATH.
-			(ex: "bin" if executables are in the bin folder)
-		default_global (str):  The default global installation path.
-			(Default is "C:\\Program Files" which is the most common location for executables on Windows)
-
-	Returns:
-		str: The installation path.
-	"""
-	# Default path is located in the current working directory
-	default_local_path: str = clean_path(os.path.join(os.getcwd(), program_name))
-
-	# Define default global path (used in prompt even if not chosen initially)
-	default_global_path: str = clean_path(os.path.join(default_global, program_name))
-
-	# Ask user for installation type (global/local)
-	install_type: str = ask_install_type(ask_global, default_local_path, default_global_path)
-
-	# If the user wants to install globally,
-	if install_type == 'g':
-
-		# Check if the user has admin privileges,
-		if not check_admin_windows():
-
-			# If the user doesn't have admin privileges, fallback to local
-			warning(
-				f"Global installation requires administrator privileges. Please re-run as administrator.\n"
-				f"Install locally instead to '{default_local_path}'? (Y/n): "
-			)
-			if input().lower() == 'n':
-				info("Installation cancelled.")
-				return ""
-			else:
-				# Fallback to local path if user agrees
-				return prompt_for_path(
-					f"Falling back to local installation path: {default_local_path}.",
-					default_local_path
-				)
-
-		# If the user has admin privileges,
-		else:
-			# Ask it user wants to override the default global install path
-			install_path: str = prompt_for_path(
-				f"Default global installation path is {default_global_path}.",
-				default_global_path
-			)
-			if add_path:
-				add_to_path_windows(os.path.join(install_path, append_to_path))
-			return install_path
-
-	# Local install
-	else: # install_type == 'l'
-		return prompt_for_path(
-			f"Default local installation path is {default_local_path}.",
-			default_local_path
-		)
-
+""" Installer module for Windows specific functions.
+
+Provides Windows specific implementations for checking administrator privileges,
+determining appropriate installation paths (global/local), and modifying
+the user's PATH environment variable.
+"""
+# Imports
+import os
+
+from ..decorators import LogLevels, handle_error
+from ..io import clean_path
+from ..print import debug, info, warning
+from .common import ask_install_type, prompt_for_path
+
+
+# Functions
+@handle_error(message="Failed to add to PATH (Windows)", error_log=LogLevels.WARNING_TRACEBACK)
+def add_to_path_windows(install_path: str) -> bool | None:
+	""" Add install_path to the User PATH environment variable on Windows.
+
+	Args:
+		install_path (str): The path to add to the User PATH environment variable.
+
+	Returns:
+		bool | None: True if the path was added to the User PATH environment variable, None otherwise.
+	"""
+	# Convert install_path to a Windows path if it's not already
+	install_path = install_path.replace("/", "\\")
+	os.makedirs(install_path, exist_ok=True)
+
+	# Get current user PATH
+	import winreg
+	with winreg.OpenKey(winreg.HKEY_CURRENT_USER, "Environment", 0, winreg.KEY_READ | winreg.KEY_WRITE) as key:
+
+		# Get the number of values in the registry key
+		num_values = winreg.QueryInfoKey(key)[1]
+
+		# Find the index of the 'Path' value
+		path_index = -1
+		for i in range(num_values):
+			if winreg.EnumValue(key, i)[0] == 'Path':
+				path_index = i
+				break
+
+		# Get the current path value
+		current_path: str = winreg.EnumValue(key, path_index)[1]
+
+		# Check if path is already present
+		if install_path not in current_path.split(';'):
+			new_path: str = f"{current_path};{install_path}"
+			winreg.SetValueEx(key, "Path", 0, winreg.REG_EXPAND_SZ, new_path)
+			debug(f"Added '{install_path}' to user PATH. Please restart your terminal for changes to take effect.")
+		else:
+			debug(f"'{install_path}' is already in user PATH.")
+	return True
+
+
+def check_admin_windows() -> bool:
+	""" Check if the script is running with administrator privileges on Windows. """
+	try:
+		import ctypes
+		return ctypes.windll.shell32.IsUserAnAdmin() != 0
+	except Exception:
+		return False
+
+
+@handle_error(message="Failed to get installation path (Windows)", error_log=LogLevels.ERROR_TRACEBACK)
+def get_install_path_windows(
+	program_name: str,
+	ask_global: int = 0,
+	add_path: bool = True,
+	append_to_path: str = "",
+	default_global: str = os.environ.get("ProgramFiles", "C:\\Program Files")
+) -> str:
+	""" Get the installation path for the program
+
+	Args:
+		program_name  (str):   The name of the program to install.
+		ask_global    (int):   0 = ask for anything, 1 = install globally, 2 = install locally
+		add_path      (bool):  Whether to add the program to the PATH environment variable. (Only if installed globally)
+		append_to_path (str):  String to append to the installation path when adding to PATH.
+			(ex: "bin" if executables are in the bin folder)
+		default_global (str):  The default global installation path.
+			(Default is "C:\\Program Files" which is the most common location for executables on Windows)
+
+	Returns:
+		str: The installation path.
+	"""
+	# Default path is located in the current working directory
+	default_local_path: str = clean_path(os.path.join(os.getcwd(), program_name))
+
+	# Define default global path (used in prompt even if not chosen initially)
+	default_global_path: str = clean_path(os.path.join(default_global, program_name))
+
+	# Ask user for installation type (global/local)
+	install_type: str = ask_install_type(ask_global, default_local_path, default_global_path)
+
+	# If the user wants to install globally,
+	if install_type == 'g':
+
+		# Check if the user has admin privileges,
+		if not check_admin_windows():
+
+			# If the user doesn't have admin privileges, fallback to local
+			warning(
+				f"Global installation requires administrator privileges. Please re-run as administrator.\n"
+				f"Install locally instead to '{default_local_path}'? (Y/n): "
+			)
+			if input().lower() == 'n':
+				info("Installation cancelled.")
+				return ""
+			else:
+				# Fallback to local path if user agrees
+				return prompt_for_path(
+					f"Falling back to local installation path: {default_local_path}.",
+					default_local_path
+				)
+
+		# If the user has admin privileges,
+		else:
+			# Ask it user wants to override the default global install path
+			install_path: str = prompt_for_path(
+				f"Default global installation path is {default_global_path}.",
+				default_global_path
+			)
+			if add_path:
+				add_to_path_windows(os.path.join(install_path, append_to_path))
+			return install_path
+
+	# Local install
+	else: # install_type == 'l'
+		return prompt_for_path(
+			f"Default local installation path is {default_local_path}.",
+			default_local_path
+		)
+

stouputils/io.py

@@ -188,16 +188,23 @@ def csv_dump(
 	done: bool = False
 
 	# Handle Polars DataFrame
-	try:
-		import polars as pl  # type: ignore
-		if isinstance(data, pl.DataFrame):
-			copy_kwargs = kwargs.copy()
-			copy_kwargs.setdefault("separator", delimiter)
-			copy_kwargs.setdefault("include_header", has_header)
-			data.write_csv(output, *args, **copy_kwargs)
-			done = True
-	except Exception:
+	import sys
+	if sys.version_info >= (3, 14) and not sys._is_gil_enabled(): # pyright: ignore[reportPrivateUsage]
+		# Skip Polars on free-threaded Python 3.14 due to segfault
+		# TODO: Remove this check when Polars is fixed
+		# See https://github.com/pola-rs/polars/issues/21889 and https://github.com/durandtibo/coola/issues/1066
 		pass
+	else:
+		try:
+			import polars as pl  # type: ignore
+			if isinstance(data, pl.DataFrame):
+				copy_kwargs = kwargs.copy()
+				copy_kwargs.setdefault("separator", delimiter)
+				copy_kwargs.setdefault("include_header", has_header)
+				data.write_csv(output, *args, **copy_kwargs)
+				done = True
+		except Exception:
+			pass
 
 	# Handle pandas DataFrame
 	if not done:

stouputils/parallel.py

@@ -7,6 +7,16 @@ This module provides utility functions for parallel processing, such as:
 
 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
 """
@@ -42,6 +52,7 @@ def multiprocessing[T, R](
 	desc: str = "",
 	max_workers: int | float = CPU_COUNT,
 	delay_first_calls: float = 0,
+	nice: int | None = None,
 	color: str = MAGENTA,
 	bar_format: str = BAR_FORMAT,
 	ascii: bool = False,
@@ -69,6 +80,11 @@ def multiprocessing[T, R](
 		delay_first_calls	(float):			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.
+		nice				(int | None):		Adjust the priority of worker processes (Defaults to None).
+			Use Unix-style values: -20 (highest priority) to 19 (lowest priority).
+			Positive values reduce priority, negative values increase it.
+			Automatically converted to appropriate priority class on Windows.
+			If None, no priority adjustment is made.
 		color				(str):				Color of the progress bar (Defaults to MAGENTA)
 		bar_format			(str):				Format of the progress bar (Defaults to BAR_FORMAT)
 		ascii				(bool):				Whether to use ASCII or Unicode characters for the progress bar
@@ -140,13 +156,21 @@ def multiprocessing[T, R](
 	# Do multiprocessing only if there is more than 1 argument and more than 1 CPU
 	if max_workers > 1 and len(args) > 1:
 		def process() -> list[Any]:
+			# Wrap function with nice if specified
+			if nice is not None:
+				wrapped_args = [(nice, func, arg) for arg in args]
+				wrapped_func = _nice_wrapper
+			else:
+				wrapped_args = args
+				wrapped_func = func
+
 			if verbose:
 				return list(process_map(
-					func, args, max_workers=max_workers, chunksize=chunksize, desc=desc, bar_format=bar_format, ascii=ascii, **tqdm_kwargs
+					wrapped_func, wrapped_args, max_workers=max_workers, chunksize=chunksize, desc=desc, bar_format=bar_format, ascii=ascii, **tqdm_kwargs
 				)) # type: ignore
 			else:
 				with Pool(max_workers) as pool:
-					return list(pool.map(func, args, chunksize=chunksize))	# type: ignore
+					return list(pool.map(wrapped_func, wrapped_args, chunksize=chunksize))	# type: ignore
 		try:
 			return process()
 		except RuntimeError as e:
@@ -382,6 +406,68 @@ def run_in_subprocess[R](
 		raise RuntimeError("Subprocess did not return any result") from e
 
 
+# "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
+			try:
+				import psutil
+				if nice_value <= -10:
+					priority = psutil.HIGH_PRIORITY_CLASS
+				elif nice_value < 0:
+					priority = psutil.ABOVE_NORMAL_PRIORITY_CLASS
+				elif nice_value == 0:
+					priority = psutil.NORMAL_PRIORITY_CLASS
+				elif nice_value < 10:
+					priority = psutil.BELOW_NORMAL_PRIORITY_CLASS
+				else:
+					priority = psutil.IDLE_PRIORITY_CLASS
+				psutil.Process().nice(priority)
+			except ImportError:
+				# Fallback to ctypes if psutil is not available
+				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 for subprocess wrapper (must be at module level for pickling on Windows)
 def _subprocess_wrapper[R](
 	result_queue: Any,

stouputils/parallel.pyi

@@ -10,7 +10,7 @@ CPU_COUNT: int
 T = TypeVar('T')
 R = TypeVar('R')
 
-def multiprocessing[T, R](func: Callable[..., R] | list[Callable[..., R]], args: Iterable[T], use_starmap: bool = False, chunksize: int = 1, desc: str = '', max_workers: int | float = ..., delay_first_calls: float = 0, color: str = ..., bar_format: str = ..., ascii: bool = False, smooth_tqdm: bool = True, **tqdm_kwargs: Any) -> list[R]:
+def multiprocessing[T, R](func: Callable[..., R] | list[Callable[..., R]], args: Iterable[T], use_starmap: bool = False, chunksize: int = 1, desc: str = '', max_workers: int | float = ..., delay_first_calls: float = 0, nice: int | None = None, color: str = ..., bar_format: str = ..., ascii: bool = False, smooth_tqdm: bool = True, **tqdm_kwargs: Any) -> 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.
@@ -32,6 +32,11 @@ def multiprocessing[T, R](func: Callable[..., R] | list[Callable[..., R]], args:
 \t\tdelay_first_calls\t(float):\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.
 \t\t\t(Defaults to 0): This can be useful to avoid functions being called in the same second.
+\t\tnice\t\t\t\t(int | None):\t\tAdjust the priority of worker processes (Defaults to None).
+\t\t\tUse Unix-style values: -20 (highest priority) to 19 (lowest priority).
+\t\t\tPositive values reduce priority, negative values increase it.
+\t\t\tAutomatically converted to appropriate priority class on Windows.
+\t\t\tIf None, no priority adjustment is made.
 \t\tcolor\t\t\t\t(str):\t\t\t\tColor of the progress bar (Defaults to MAGENTA)
 \t\tbar_format\t\t\t(str):\t\t\t\tFormat of the progress bar (Defaults to BAR_FORMAT)
 \t\tascii\t\t\t\t(bool):\t\t\t\tWhether to use ASCII or Unicode characters for the progress bar
@@ -169,6 +174,21 @@ def run_in_subprocess[R](func: Callable[..., R], *args: Any, timeout: float | No
 \t\t\t> # With timeout to prevent hanging
 \t\t\t> run_in_subprocess(some_gpu_func, data, timeout=300.0)
 \t'''
+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 _subprocess_wrapper[R](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.