stouputils.decorators — stouputils 1.0.0 documentation

# Source code for stouputils.decorators

    """
    This module provides decorators for various purposes:

    - silent(): Make a function silent (disable stdout, and stderr if specified) (alternative to stouputils.ctx.Muffle)
    - 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
    - deprecated(): Mark a function as deprecated
    """

    # Imports
    import os
    import sys
    import time
    from enum import Enum
    from pickle import dumps as pickle_dumps
    from traceback import format_exc
    from typing import Callable, Literal, Any
    from functools import wraps
    from .print import *

    # Decorator that make a function silent (disable stdout)

    def silent(
    	func: Callable[..., Any],
    	mute_stderr: bool = False
    ) -> Callable[..., Any]:
    	""" Decorator that makes a function silent (disable stdout, and stderr if specified).

    	Alternative to stouputils.ctx.Muffle.

    	Args:
    		func			(Callable[..., Any]):	Function to make silent
    		mute_stderr		(bool):					Whether to mute stderr or not

    	Examples:
    		>>> @silent
    		... def test():
    		...     print("Hello, world!")
    		>>> test()

    		>>> silent(print)("Hello, world!")
    	"""
    	@wraps(func)
    	def wrapper(*args: tuple[Any, ...], **kwargs: dict[str, Any]) -> Any:

    		# Disable stdout and stderr
    		_original_stdout: Any = sys.stdout
    		_original_stderr: Any = None
    		sys.stdout = open(os.devnull, "w", encoding="utf-8")
    		if mute_stderr:
    			_original_stderr = sys.stderr
    			sys.stderr = open(os.devnull, "w", encoding="utf-8")

    		# Call the function
    		result: Any = func(*args, **kwargs)

    		# Re-Enable stdout and stderr
    		sys.stdout.close()
    		sys.stdout = _original_stdout
    		if mute_stderr:
    			sys.stderr.close()
    			sys.stderr = _original_stderr
    		return result
    	return wrapper

    # Execution time decorator

    def measure_time(
    	print_func: Callable[..., None] = debug,
    	message: str = "",
    	perf_counter: bool = True
    ) -> Callable[..., Any]:
    	""" Decorator that will measure the execution time of a function and print it with the given 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 (e.g. "Execution time of Something"), defaults to "Execution time of {func.__name__}"
    		perf_counter	(bool):		Whether to use time.perf_counter_ns or time.time_ns
    	Returns:
    		Callable:	Decorator to measure the time of the function.

    	Examples:
    		>>> @measure_time(info)
    		... def test():
    		...     pass
    		>>> test()	# [INFO HH:MM:SS] Execution time of test: 0.000ms (0ns)
    	"""
    	ns: Callable[[], int] = time.perf_counter_ns if perf_counter else time.time_ns
    	def decorator(func: Callable[..., Any]) -> Callable[..., Any]:

    		# Set the message if not specified
    		nonlocal message
    		if not message:
    			message = f"Execution time of {func.__name__}"

    		@wraps(func)
    		def wrapper(*args: tuple[Any, ...], **kwargs: dict[str, Any]) -> Any:

    			# Measure the execution time (nanoseconds and seconds)
    			start_ns: int = ns()
    			result = func(*args, **kwargs)
    			total_ns: int = ns() - 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.3s, seconds otherwise)
    			if total_ms < 300:
    				print_func(f"{message}: {total_ms:.3f}ms ({total_ns}ns)")
    			elif total_s < 60:
    				print_func(f"{message}: {(total_s):.5f}s")
    			else:
    				minutes: int = int(total_s) // 60
    				seconds: int = int(total_s) % 60
    				if minutes < 60:
    					print_func(f"{message}: {minutes}m {seconds}s")
    				else:
    					hours: int = minutes // 60
    					minutes: int = minutes % 60
    					if hours < 24:
    						print_func(f"{message}: {hours}h {minutes}m {seconds}s")
    					else:
    						days: int = hours // 24
    						hours: int = hours % 24
    						print_func(f"{message}: {days}d {hours}h {minutes}m {seconds}s")
    			return result
    		return wrapper
    	return decorator

    # Decorator that handle an error with different log levels

    class LogLevels(Enum):
    	""" Log level for the errors in the decorator handle_error() """
    	NONE = 0
    	""" Do nothing """
    	WARNING = 1
    	""" Show as warning """
    	WARNING_TRACEBACK = 2
    	""" Show as warning with traceback """
    	ERROR_TRACEBACK = 3
    	""" Show as error with traceback """
    	RAISE_EXCEPTION = 4
    	""" Raise exception """

    force_raise_exception: bool = False
    """ If true, the error_log parameter will be set to RAISE_EXCEPTION for every next handle_error calls, useful for doctests """

    def handle_error(
    	exceptions: tuple[type[BaseException], ...] | type[BaseException] = (Exception,),
    	message: str = "",
    	error_log: LogLevels = LogLevels.ERROR_TRACEBACK
    ) -> Callable[..., Any]:
    	""" Decorator that handle an error with different log levels.

    	Args:
    		exceptions		(tuple[type[BaseException]], ...):	Exceptions to handle
    		message			(str):								Message to display with the error. (e.g. "Error during something")
    		error_log		(LogLevels):						Log level for the errors
    			LogLevels.NONE:					None
    			LogLevels.WARNING:				Show as warning
    			LogLevels.WARNING_TRACEBACK:	Show as warning with traceback
    			LogLevels.ERROR_TRACEBACK:		Show as error with traceback
    			LogLevels.RAISE_EXCEPTION:		Raise exception (as if the decorator didn't exist)

    	Examples:
    		>>> @handle_error(error_log=LogLevels.WARNING)
    		... def test():
    		...     raise ValueError("Let's fail")
    		>>> test()	# [WARNING HH:MM:SS] Error during test: (ValueError) Let's fail
    	"""
    	# Update error_log if needed
    	if force_raise_exception:
    		error_log = LogLevels.RAISE_EXCEPTION

    	def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
    		if message != "":
    			msg: str = f"{message}, "
    		else:
    			msg: str = message

    		@wraps(func)
    		def wrapper(*args: tuple[Any, ...], **kwargs: dict[str, Any]) -> Any:
    			try:
    				return func(*args, **kwargs)
    			except exceptions as e:
    				if error_log == LogLevels.WARNING:
    					warning(f"{msg}Error during {func.__name__}: ({type(e).__name__}) {e}")
    				elif error_log == LogLevels.WARNING_TRACEBACK:
    					warning(f"{msg}Error during {func.__name__}:\n{format_exc()}")
    				elif error_log == LogLevels.ERROR_TRACEBACK:
    					error(f"{msg}Error during {func.__name__}:\n{format_exc()}", exit=True)
    				elif error_log == LogLevels.RAISE_EXCEPTION:
    					raise e
    		return wrapper
    	return decorator

    # Easy cache function with parameter caching method

    def simple_cache(method: Literal["str", "pickle"] = "str") -> Callable[..., 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).

    	Args:
    		method (Literal["str", "pickle"]): The method to use for caching.
    	Returns:
    		Callable[..., Callable[..., Any]]: A decorator that caches the result of a function.
    	Examples:
    		>>> @simple_cache(method="str")
    		... def test(a: int, b: int) -> int:
    		...     return a + b
    		>>> test(1, 2)	# 3
    		3
    		>>> test(1, 2)	# 3
    		3
    		>>> test(3, 4)	# 7
    		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
    		return wrapper

    	# Return the decorator
    	return decorator

    def deprecated(
    	message: str = "",
    	error_log: LogLevels = LogLevels.WARNING
    ) -> Callable[..., Any]:
    	""" Decorator that marks a function as deprecated.

    	Args:
    		message     (str):               Additional message to display with the deprecation warning
    		error_log   (LogLevels):         Log level for the deprecation warning
    			LogLevels.NONE:              None
    			LogLevels.WARNING:           Show as warning
    			LogLevels.WARNING_TRACEBACK: Show as warning with traceback
    			LogLevels.ERROR_TRACEBACK:   Show as error with traceback
    			LogLevels.RAISE_EXCEPTION:   Raise exception
    	Returns:
    		Callable[..., Any]: Decorator that marks a function as deprecated

    	Examples:
    		>>> @deprecated(message="Use 'this_function()' instead", error_log=LogLevels.WARNING)
    		... def test():
    		...     pass
    		>>> test()	# [WARNING HH:MM:SS] Function 'test()' is deprecated. Use 'this_function()' instead
    	"""
    	def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
    		@wraps(func)
    		def wrapper(*args: tuple[Any, ...], **kwargs: dict[str, Any]) -> Any:
    			# Build deprecation message
    			msg: str = f"Function '{func.__name__}()' is deprecated"
    			if message:
    				msg += f". {message}"

    			# Handle deprecation warning based on log level
    			if error_log == LogLevels.WARNING:
    				warning(msg)
    			elif error_log == LogLevels.WARNING_TRACEBACK:
    				warning(f"{msg}\n{format_exc()}")
    			elif error_log == LogLevels.ERROR_TRACEBACK:
    				error(f"{msg}\n{format_exc()}", exit=True)
    			elif error_log == LogLevels.RAISE_EXCEPTION:
    				raise DeprecationWarning(msg)

    			# Call the original function
    			return func(*args, **kwargs)
    		return wrapper
    	return decorator

    def __test_simple_cache():
    	@measure_time(progress)
    	@simple_cache(method="pickle")
    	def test_pickle_cache(*args: tuple[Any, ...], **kwargs: dict[str, Any]) -> str:
    		return "pickle cache"

    	@measure_time(progress)
    	@simple_cache(method="str")
    	def test_str_cache(*args: tuple[Any, ...], **kwargs: dict[str, Any]) -> str:
    		return "str cache"

    	args: list[int] = list(range(1_000_000))
    	kwargs: dict[str, int] = {str(i): i for i in args}

    	info("Testing with large arguments")
    	for _ in range(2):
    		test_pickle_cache(*args, **kwargs)
    		test_str_cache(*args, **kwargs)

    	info("Testing with small arguments")
    	for _ in range(2):
    		test_pickle_cache()
    		test_str_cache()

    if __name__ == "__main__":
    	__test_simple_cache()

    	@deprecated(message="Use 'this_function()' instead", error_log=LogLevels.WARNING)
    	def test_deprecated():
    		return "test"
    	test_deprecated()