stouputils 1.12.1__py3-none-any.whl

stouputils/ctx.py

@@ -0,0 +1,408 @@
+"""
+This module provides context managers for various utilities such as logging to a file,
+measuring execution time, silencing output, and setting multiprocessing start methods.
+
+- LogToFile: Context manager to log to a file every print call (with LINE_UP handling)
+- 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
+"""
+
+# Imports
+from __future__ import annotations
+
+import os
+import sys
+import time
+from collections.abc import Callable
+from contextlib import AbstractAsyncContextManager, AbstractContextManager
+from typing import IO, Any, TextIO, TypeVar
+
+from .io import super_open
+from .print import TeeMultiOutput, debug
+
+# Type variable for context managers
+T = TypeVar("T")
+
+# Abstract base class for context managers supporting both sync and async usage
+class AbstractBothContextManager[T](AbstractContextManager[T], AbstractAsyncContextManager[T]):
+    """ Abstract base class for context managers that support both synchronous and asynchronous usage. """
+    pass
+
+# Context manager to log to a file
+class LogToFile(AbstractBothContextManager["LogToFile"]):
+	""" Context manager to log to a file.
+
+	This context manager allows you to temporarily log output to a file while still printing normally.
+	The file will receive log messages without ANSI color codes.
+
+	Args:
+		path (str): Path to the log file
+		mode (str): Mode to open the file in (default: "w")
+		encoding (str): Encoding to use for the file (default: "utf-8")
+		tee_stdout (bool): Whether to redirect stdout to the file (default: True)
+		tee_stderr (bool): Whether to redirect stderr to the file (default: True)
+		ignore_lineup (bool): Whether to ignore lines containing LINE_UP escape sequence in files (default: False)
+		restore_on_exit (bool): Whether to restore original stdout/stderr on exit (default: False)
+			This ctx uses TeeMultiOutput which handles closed files gracefully, so restoring is not mandatory.
+
+	Examples:
+		.. code-block:: python
+
+			> import stouputils as stp
+			> with stp.LogToFile("output.log"):
+			>     stp.info("This will be logged to output.log and printed normally")
+			>     print("This will also be logged")
+
+			> with stp.LogToFile("output.log") as log_ctx:
+			>     stp.warning("This will be logged to output.log and printed normally")
+			>     log_ctx.change_file("new_file.log")
+			>     print("This will be logged to new_file.log")
+	"""
+	def __init__(
+		self,
+		path: str,
+		mode: str = "w",
+		encoding: str = "utf-8",
+		tee_stdout: bool = True,
+		tee_stderr: bool = True,
+		ignore_lineup: bool = True,
+		restore_on_exit: bool = False
+	) -> None:
+		self.path: str = path
+		""" Attribute remembering path to the log file """
+		self.mode: str = mode
+		""" Attribute remembering mode to open the file in """
+		self.encoding: str = encoding
+		""" Attribute remembering encoding to use for the file """
+		self.tee_stdout: bool = tee_stdout
+		""" Whether to redirect stdout to the file """
+		self.tee_stderr: bool = tee_stderr
+		""" Whether to redirect stderr to the file """
+		self.ignore_lineup: bool = ignore_lineup
+		""" Whether to ignore lines containing LINE_UP escape sequence in files """
+		self.restore_on_exit: bool = restore_on_exit
+		""" Whether to restore original stdout/stderr on exit.
+		This ctx uses TeeMultiOutput which handles closed files gracefully, so restoring is not mandatory. """
+		self.file: IO[Any]
+		""" Attribute remembering opened file """
+		self.original_stdout: TextIO
+		""" Original stdout before redirection """
+		self.original_stderr: TextIO
+		""" Original stderr before redirection """
+
+	def __enter__(self) -> LogToFile:
+		""" Enter context manager which opens the log file and redirects stdout/stderr """
+		# Open file
+		self.file = super_open(self.path, mode=self.mode, encoding=self.encoding)
+
+		# Redirect stdout and stderr if requested
+		if self.tee_stdout:
+			self.original_stdout = sys.stdout
+			sys.stdout = TeeMultiOutput(self.original_stdout, self.file, ignore_lineup=self.ignore_lineup)
+		if self.tee_stderr:
+			self.original_stderr = sys.stderr
+			sys.stderr = TeeMultiOutput(self.original_stderr, self.file, ignore_lineup=self.ignore_lineup)
+
+		# Return self
+		return self
+
+	def __exit__(self, exc_type: type[BaseException]|None, exc_val: BaseException|None, exc_tb: Any|None) -> None:
+		""" Exit context manager which closes the log file and restores stdout/stderr """
+		# Restore original stdout and stderr (if requested)
+		if self.restore_on_exit:
+			if self.tee_stdout:
+				sys.stdout = self.original_stdout
+			if self.tee_stderr:
+				sys.stderr = self.original_stderr
+
+		# Close file
+		self.file.close()
+
+	async def __aenter__(self) -> LogToFile:
+		""" Enter async context manager which opens the log file and redirects stdout/stderr """
+		return self.__enter__()
+
+	async def __aexit__(self, exc_type: type[BaseException]|None, exc_val: BaseException|None, exc_tb: Any|None) -> None:
+		""" Exit async context manager which closes the log file and restores stdout/stderr """
+		self.__exit__(exc_type, exc_val, exc_tb)
+
+	def change_file(self, new_path: str) -> None:
+		""" Change the log file to a new path.
+
+		Args:
+			new_path (str): New path to the log file
+		"""
+		# Close current file, open new file and redirect outputs
+		self.file.close()
+		self.path = new_path
+		self.__enter__()
+
+	@staticmethod
+	def common(logs_folder: str, filepath: str, func: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
+		""" Common code used at the beginning of a program to launch main function
+
+		Args:
+			logs_folder (str): Folder to store logs in
+			filepath    (str): Path to the main function
+			func        (Callable[..., Any]): Main function to launch
+			*args       (tuple[Any, ...]): Arguments to pass to the main function
+			**kwargs    (dict[str, Any]): Keyword arguments to pass to the main function
+		Returns:
+			Any: Return value of the main function
+
+		Examples:
+			>>> if __name__ == "__main__":
+			...     LogToFile.common(f"{ROOT}/logs", __file__, main)
+		"""
+		# Import datetime
+		from datetime import datetime
+
+		# Build log file path
+		file_basename: str = os.path.splitext(os.path.basename(filepath))[0]
+		date_time: str = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+		date_str, time_str = date_time.split("_")
+		log_filepath: str = f"{logs_folder}/{file_basename}/{date_str}/{time_str}.log"
+
+		# Launch function with arguments if any
+		with LogToFile(log_filepath):
+			return func(*args, **kwargs)
+
+# Context manager to measure execution time
+class MeasureTime(AbstractBothContextManager["MeasureTime"]):
+	""" Context manager to measure execution time.
+
+	This context manager measures the execution time of the code block it wraps
+	and prints the result using a specified print function.
+
+	Args:
+		print_func      (Callable): Function to use to print the execution time (e.g. debug, info, warning, error, etc.).
+		message         (str):      Message to display with the execution time. Defaults to "Execution time".
+		perf_counter    (bool):     Whether to use time.perf_counter_ns or time.time_ns. Defaults to True.
+
+	Examples:
+		.. code-block:: python
+
+			> import time
+			> import stouputils as stp
+			> with stp.MeasureTime(stp.info, message="My operation"):
+			...     time.sleep(0.5)
+			> # [INFO HH:MM:SS] My operation: 500.123ms (500123456ns)
+
+			> with stp.MeasureTime(): # Uses debug by default
+			...     time.sleep(0.1)
+			> # [DEBUG HH:MM:SS] Execution time: 100.456ms (100456789ns)
+	"""
+	def __init__(
+		self,
+		print_func: Callable[..., None] = debug,
+		message: str = "Execution time",
+		perf_counter: bool = True
+	) -> None:
+		self.print_func: Callable[..., None] = print_func
+		""" Function to use for printing the execution time """
+		self.message: str = message
+		""" Message to display with the execution time """
+		self.perf_counter: bool = perf_counter
+		""" Whether to use time.perf_counter_ns or time.time_ns """
+		self.ns: Callable[[], int] = time.perf_counter_ns if perf_counter else time.time_ns
+		""" Time function to use """
+		self.start_ns: int = 0
+		""" Start time in nanoseconds """
+
+	def __enter__(self) -> MeasureTime:
+		""" Enter context manager, record start time """
+		self.start_ns = self.ns()
+		return self
+
+	def __exit__(self, exc_type: type[BaseException]|None, exc_val: BaseException|None, exc_tb: Any|None) -> None:
+		""" Exit context manager, calculate duration and print """
+		# Measure the execution time (nanoseconds and seconds)
+		total_ns: int = self.ns() - self.start_ns
+		total_ms: float = total_ns / 1_000_000
+		total_s: float = total_ns / 1_000_000_000
+
+		# Print the execution time (nanoseconds if less than 0.1s, seconds otherwise)
+		if total_ms < 100:
+			self.print_func(f"{self.message}: {total_ms:.3f}ms ({total_ns}ns)")
+		elif total_s < 60:
+			self.print_func(f"{self.message}: {(total_s):.5f}s")
+		else:
+			minutes: int = int(total_s) // 60
+			seconds: int = int(total_s) % 60
+			if minutes < 60:
+				self.print_func(f"{self.message}: {minutes}m {seconds}s")
+			else:
+				hours: int = minutes // 60
+				minutes: int = minutes % 60
+				if hours < 24:
+					self.print_func(f"{self.message}: {hours}h {minutes}m {seconds}s")
+				else:
+					days: int = hours // 24
+					hours: int = hours % 24
+					self.print_func(f"{self.message}: {days}d {hours}h {minutes}m {seconds}s")
+
+	async def __aenter__(self) -> MeasureTime:
+		""" Enter async context manager, record start time """
+		return self.__enter__()
+
+	async def __aexit__(self, exc_type: type[BaseException]|None, exc_val: BaseException|None, exc_tb: Any|None) -> None:
+		""" Exit async context manager, calculate duration and print """
+		self.__exit__(exc_type, exc_val, exc_tb)
+
+# Context manager to temporarily silence output
+class Muffle(AbstractBothContextManager["Muffle"]):
+	""" Context manager that temporarily silences output.
+	(No thread-safety guaranteed)
+
+	Alternative to stouputils.decorators.silent()
+
+	Examples:
+		>>> with Muffle():
+		...     print("This will not be printed")
+	"""
+	def __init__(self, mute_stderr: bool = False) -> None:
+		self.mute_stderr: bool = mute_stderr
+		""" Attribute remembering if stderr should be muted """
+		self.original_stdout: IO[Any]
+		""" Attribute remembering original stdout """
+		self.original_stderr: IO[Any]
+		""" Attribute remembering original stderr """
+
+	def __enter__(self) -> Muffle:
+		""" Enter context manager which redirects stdout and stderr to devnull """
+		# Redirect stdout to devnull
+		self.original_stdout = sys.stdout
+		sys.stdout = open(os.devnull, "w", encoding="utf-8")
+
+		# Redirect stderr to devnull if needed
+		if self.mute_stderr:
+			self.original_stderr = sys.stderr
+			sys.stderr = open(os.devnull, "w", encoding="utf-8")
+
+		# Return self
+		return self
+
+	def __exit__(self, exc_type: type[BaseException]|None, exc_val: BaseException|None, exc_tb: Any|None) -> None:
+		""" Exit context manager which restores original stdout and stderr """
+		# Restore original stdout
+		sys.stdout.close()
+		sys.stdout = self.original_stdout
+
+		# Restore original stderr if needed
+		if self.mute_stderr:
+			sys.stderr.close()
+			sys.stderr = self.original_stderr
+
+	async def __aenter__(self) -> Muffle:
+		""" Enter async context manager which redirects stdout and stderr to devnull """
+		return self.__enter__()
+
+	async def __aexit__(self, exc_type: type[BaseException]|None, exc_val: BaseException|None, exc_tb: Any|None) -> None:
+		""" Exit async context manager which restores original stdout and stderr """
+		self.__exit__(exc_type, exc_val, exc_tb)
+
+# Context manager that does nothing
+class DoNothing(AbstractBothContextManager["DoNothing"]):
+	""" Context manager that does nothing.
+
+	This is a no-op context manager that can be used as a placeholder
+	or for conditional context management.
+
+	Different from contextlib.nullcontext because it handles args and kwargs,
+	along with **async** context management.
+
+	Examples:
+		>>> with DoNothing():
+		...     print("This will be printed normally")
+		This will be printed normally
+
+		>>> # Conditional context management
+		>>> some_condition = True
+		>>> ctx = DoNothing() if some_condition else Muffle()
+		>>> with ctx:
+		...     print("May or may not be printed depending on condition")
+		May or may not be printed depending on condition
+	"""
+	def __init__(self, *args: Any, **kwargs: Any) -> None:
+		""" No initialization needed, this is a no-op context manager """
+		pass
+
+	def __enter__(self) -> DoNothing:
+		""" Enter context manager (does nothing) """
+		return self
+
+	def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+		""" Exit context manager (does nothing) """
+		pass
+
+	async def __aenter__(self) -> DoNothing:
+		""" Enter async context manager (does nothing) """
+		return self
+
+	async def __aexit__(self, *excinfo: Any) -> None:
+		""" Exit async context manager (does nothing) """
+		pass
+NullContextManager = DoNothing
+""" Alias for DoNothing context manager """
+
+# Context manager to temporarily set multiprocessing start method
+class SetMPStartMethod(AbstractBothContextManager["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)
+
+	async def __aenter__(self) -> SetMPStartMethod:
+		""" Enter async context manager which sets the start method """
+		return self.__enter__()
+
+	async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+		""" Exit async context manager which restores the original start method """
+		self.__exit__(exc_type, exc_val, exc_tb)
+

stouputils/ctx.pyi

@@ -0,0 +1,211 @@
+import abc
+from .io import super_open as super_open
+from .print import TeeMultiOutput as TeeMultiOutput, debug as debug
+from collections.abc import Callable as Callable
+from contextlib import AbstractAsyncContextManager, AbstractContextManager
+from typing import Any, IO, TextIO, TypeVar
+
+T = TypeVar('T')
+
+class AbstractBothContextManager[T](AbstractContextManager[T], AbstractAsyncContextManager[T], metaclass=abc.ABCMeta):
+    """ Abstract base class for context managers that support both synchronous and asynchronous usage. """
+
+class LogToFile(AbstractBothContextManager['LogToFile']):
+    ''' Context manager to log to a file.
+
+\tThis context manager allows you to temporarily log output to a file while still printing normally.
+\tThe file will receive log messages without ANSI color codes.
+
+\tArgs:
+\t\tpath (str): Path to the log file
+\t\tmode (str): Mode to open the file in (default: "w")
+\t\tencoding (str): Encoding to use for the file (default: "utf-8")
+\t\ttee_stdout (bool): Whether to redirect stdout to the file (default: True)
+\t\ttee_stderr (bool): Whether to redirect stderr to the file (default: True)
+\t\tignore_lineup (bool): Whether to ignore lines containing LINE_UP escape sequence in files (default: False)
+\t\trestore_on_exit (bool): Whether to restore original stdout/stderr on exit (default: False)
+\t\t\tThis ctx uses TeeMultiOutput which handles closed files gracefully, so restoring is not mandatory.
+
+\tExamples:
+\t\t.. code-block:: python
+
+\t\t\t> import stouputils as stp
+\t\t\t> with stp.LogToFile("output.log"):
+\t\t\t>     stp.info("This will be logged to output.log and printed normally")
+\t\t\t>     print("This will also be logged")
+
+\t\t\t> with stp.LogToFile("output.log") as log_ctx:
+\t\t\t>     stp.warning("This will be logged to output.log and printed normally")
+\t\t\t>     log_ctx.change_file("new_file.log")
+\t\t\t>     print("This will be logged to new_file.log")
+\t'''
+    path: str
+    mode: str
+    encoding: str
+    tee_stdout: bool
+    tee_stderr: bool
+    ignore_lineup: bool
+    restore_on_exit: bool
+    file: IO[Any]
+    original_stdout: TextIO
+    original_stderr: TextIO
+    def __init__(self, path: str, mode: str = 'w', encoding: str = 'utf-8', tee_stdout: bool = True, tee_stderr: bool = True, ignore_lineup: bool = True, restore_on_exit: bool = False) -> None: ...
+    def __enter__(self) -> LogToFile:
+        """ Enter context manager which opens the log file and redirects stdout/stderr """
+    def __exit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any | None) -> None:
+        """ Exit context manager which closes the log file and restores stdout/stderr """
+    async def __aenter__(self) -> LogToFile:
+        """ Enter async context manager which opens the log file and redirects stdout/stderr """
+    async def __aexit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any | None) -> None:
+        """ Exit async context manager which closes the log file and restores stdout/stderr """
+    def change_file(self, new_path: str) -> None:
+        """ Change the log file to a new path.
+
+\t\tArgs:
+\t\t\tnew_path (str): New path to the log file
+\t\t"""
+    @staticmethod
+    def common(logs_folder: str, filepath: str, func: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
+        ''' Common code used at the beginning of a program to launch main function
+
+\t\tArgs:
+\t\t\tlogs_folder (str): Folder to store logs in
+\t\t\tfilepath    (str): Path to the main function
+\t\t\tfunc        (Callable[..., Any]): Main function to launch
+\t\t\t*args       (tuple[Any, ...]): Arguments to pass to the main function
+\t\t\t**kwargs    (dict[str, Any]): Keyword arguments to pass to the main function
+\t\tReturns:
+\t\t\tAny: Return value of the main function
+
+\t\tExamples:
+\t\t\t>>> if __name__ == "__main__":
+\t\t\t...     LogToFile.common(f"{ROOT}/logs", __file__, main)
+\t\t'''
+
+class MeasureTime(AbstractBothContextManager['MeasureTime']):
+    ''' Context manager to measure execution time.
+
+\tThis context manager measures the execution time of the code block it wraps
+\tand prints the result using a specified print function.
+
+\tArgs:
+\t\tprint_func      (Callable): Function to use to print the execution time (e.g. debug, info, warning, error, etc.).
+\t\tmessage         (str):      Message to display with the execution time. Defaults to "Execution time".
+\t\tperf_counter    (bool):     Whether to use time.perf_counter_ns or time.time_ns. Defaults to True.
+
+\tExamples:
+\t\t.. code-block:: python
+
+\t\t\t> import time
+\t\t\t> import stouputils as stp
+\t\t\t> with stp.MeasureTime(stp.info, message="My operation"):
+\t\t\t...     time.sleep(0.5)
+\t\t\t> # [INFO HH:MM:SS] My operation: 500.123ms (500123456ns)
+
+\t\t\t> with stp.MeasureTime(): # Uses debug by default
+\t\t\t...     time.sleep(0.1)
+\t\t\t> # [DEBUG HH:MM:SS] Execution time: 100.456ms (100456789ns)
+\t'''
+    print_func: Callable[..., None]
+    message: str
+    perf_counter: bool
+    ns: Callable[[], int]
+    start_ns: int
+    def __init__(self, print_func: Callable[..., None] = ..., message: str = 'Execution time', perf_counter: bool = True) -> None: ...
+    def __enter__(self) -> MeasureTime:
+        """ Enter context manager, record start time """
+    def __exit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any | None) -> None:
+        """ Exit context manager, calculate duration and print """
+    async def __aenter__(self) -> MeasureTime:
+        """ Enter async context manager, record start time """
+    async def __aexit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any | None) -> None:
+        """ Exit async context manager, calculate duration and print """
+
+class Muffle(AbstractBothContextManager['Muffle']):
+    ''' Context manager that temporarily silences output.
+\t(No thread-safety guaranteed)
+
+\tAlternative to stouputils.decorators.silent()
+
+\tExamples:
+\t\t>>> with Muffle():
+\t\t...     print("This will not be printed")
+\t'''
+    mute_stderr: bool
+    original_stdout: IO[Any]
+    original_stderr: IO[Any]
+    def __init__(self, mute_stderr: bool = False) -> None: ...
+    def __enter__(self) -> Muffle:
+        """ Enter context manager which redirects stdout and stderr to devnull """
+    def __exit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any | None) -> None:
+        """ Exit context manager which restores original stdout and stderr """
+    async def __aenter__(self) -> Muffle:
+        """ Enter async context manager which redirects stdout and stderr to devnull """
+    async def __aexit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any | None) -> None:
+        """ Exit async context manager which restores original stdout and stderr """
+
+class DoNothing(AbstractBothContextManager['DoNothing']):
+    ''' Context manager that does nothing.
+
+\tThis is a no-op context manager that can be used as a placeholder
+\tor for conditional context management.
+
+\tDifferent from contextlib.nullcontext because it handles args and kwargs,
+\talong with **async** context management.
+
+\tExamples:
+\t\t>>> with DoNothing():
+\t\t...     print("This will be printed normally")
+\t\tThis will be printed normally
+
+\t\t>>> # Conditional context management
+\t\t>>> some_condition = True
+\t\t>>> ctx = DoNothing() if some_condition else Muffle()
+\t\t>>> with ctx:
+\t\t...     print("May or may not be printed depending on condition")
+\t\tMay or may not be printed depending on condition
+\t'''
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        """ No initialization needed, this is a no-op context manager """
+    def __enter__(self) -> DoNothing:
+        """ Enter context manager (does nothing) """
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        """ Exit context manager (does nothing) """
+    async def __aenter__(self) -> DoNothing:
+        """ Enter async context manager (does nothing) """
+    async def __aexit__(self, *excinfo: Any) -> None:
+        """ Exit async context manager (does nothing) """
+NullContextManager = DoNothing
+
+class SetMPStartMethod(AbstractBothContextManager['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 """
+    async def __aenter__(self) -> SetMPStartMethod:
+        """ Enter async context manager which sets the start method """
+    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        """ Exit async context manager which restores the original start method """

stouputils/data_science/config/get.py

@@ -0,0 +1,51 @@
+""" Load configuration from the set.py file and handle some special cases.
+
+Proper way to get the configuration is by importing this module, not the set.py file directly.
+"""
+
+# pyright: reportUnknownMemberType=false
+# pyright: reportUnknownVariableType=false
+# pyright: reportMissingTypeStubs=false
+
+# Imports
+import os
+from typing import Any
+
+from .set import DataScienceConfig
+
+# Special cases
+# Hide GPU when using CPU
+if DataScienceConfig.TENSORFLOW_DEVICE.lower().startswith("/cpu"):
+	os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
+
+# Precise which GPU we use
+elif DataScienceConfig.TENSORFLOW_DEVICE.lower().startswith("/gpu"):
+	os.environ["CUDA_VISIBLE_DEVICES"] = DataScienceConfig.TENSORFLOW_DEVICE.split(":")[-1]
+
+	# Configure TensorFlow (if available)
+	try:
+		from tensorflow import config as tf_config
+
+		# Get the physical devices
+		physical_devices: list[Any] = tf_config.list_physical_devices("GPU")
+
+		# Configure TensorFlow GPU memory management to allocate memory dynamically
+		# This prevents TensorFlow from allocating all GPU memory upfront
+		# Instead, memory will grow as needed, allowing better resource sharing
+		for device in physical_devices:
+			tf_config.experimental.set_memory_growth(device, True)
+
+		# Disable eager execution mode in TensorFlow
+		# This improves performance by allowing TensorFlow to create an optimized graph
+		# of operations instead of executing operations one by one (at the cost of debugging difficulty)
+		tf_config.run_functions_eagerly(False)
+	except ImportError:
+		pass
+
+	# Enable mixed precision training (if available)
+	try:
+		from keras import mixed_precision
+		mixed_precision.set_global_policy(DataScienceConfig.MIXED_PRECISION_POLICY)
+	except ImportError:
+		pass
+