Added Exception handeling to logwrap

Author: HMohammad2520

classmods/_decorators.py

@@ -2,67 +2,116 @@
 from functools import wraps
 from typing import Any, Callable, Literal, Optional, Tuple, TypeAlias
 
-LOG_LEVEL: TypeAlias = Literal['CRITICAL','ERROR','WARNING','INFO','DEBUG','NOTSET','DEFAULT']
+LOG_LEVEL: TypeAlias = Literal['CRITICAL', 'ERROR', 'WARNING', 'INFO', 'DEBUG', 'NOTSET', 'DEFAULT']
 
 def logwrap(
-        before: Optional[Tuple[LOG_LEVEL, str]] | str = None,
-        after: Optional[Tuple[LOG_LEVEL, str]] | str = None,
-        show_args: Optional[LOG_LEVEL] = None
-        ) -> Callable:
+        before: Optional[Tuple[LOG_LEVEL, str]] | str | bool = None,
+        on_exception: Optional[Tuple[LOG_LEVEL, str]] | str | bool = None,
+        after: Optional[Tuple[LOG_LEVEL, str]] | str | bool = None,
+    ) -> Callable:
     """
     Simple dynamic decorator to log function calls. Uses the logging module with your current project configurations.
-    Use LOG_LEVEL literal for using standard log levels. 
-    **Warning**: Providing wrong log level won't raise any exception but will default to specified defaults with DEFAULT level.
+    Use LOG_LEVEL literal for using standard log levels.
+    The messages will get formated in proccess so you can use templating for better logging.
+
+    (`func`: Function Name, `args`: Arguments Tuple, `kwargs`: Keyword Arguments Dict, `e`: The Exception if Exists)
 
+    Passing bool will use default levels and messages.(
+    Before: 'DEBUG - Calling {func} - args:{args}, kwargs:{kwargs}'
+    After: 'INFO - Function {func} ended'
+    On Exception: 'ERROR - Error on {func}: {e}')
+    
+    **Warning**: If options are negative, it will not log the option.
+    **Warning**: Providing wrong log level won't raise any exception but will default to specified defaults with DEFAULT level.
+        
     Args:
         before: Tuple of log level and message to log before function call or string with default of `DEBUG` level.
+        on_exception: Tuple of log level and message to log exception of function call or string with default of `ERROR` level.
         after: Tuple of log level and message to log after function call or string with default of `INFO` level.
-        show_args: Log level to use for logging function parameters before a call (eg. Args and kwargs) defaults to `DEBUG` level.
 
     Examples:
         >>> # Example with Custom Levels
-        >>> @logwrap(before=('INFO', 'Function starting'), after=('INFO', 'Function ended'), show_args='DEBUG')
+        >>> @logwrap(before=('INFO', '{func} starting, args={args} kwargs={kwargs}'), after=('INFO', '{my_func} ended'))
         >>> def my_func(my_arg, my_kwarg=None):
         >>>     ...
         >>> my_func('hello', my_kwarg=123) # calling the function
-        Info - My Function is Starting
-        Debug - Arguments: args=('hello',), kwargs={'my_kwarg': 123}
-        Info - Function Ended
+        Info - my_func Starting, args=('hello',), kwargs={'my_kwarg': 123}
+        Info - my_func Ended
         >>> # Example with defaults.
-        >>> @logwrap(before='Starting', after='Ended')
+        >>> @logwrap(before=True, after=True)
         >>> def my_new_func():
         >>>     ...
         >>> my_new_func()
-        Debug - Starting
-        Info - Ended
+        Debug - Calling my_new_func - args:(), kwargs:{}
+        Info - Function my_new_func ended
+        >>> # Example with Exception
+        >>> @logwrap(on_exception=True)
+        >>> def error_func():
+        >>>     raise Exception('My exception msg')
+        >>> error_func()
+        Error - Error on error_func: My exception msg')
     """
-    if isinstance(before, str):
-        before = ('DEFAULT', before)
+    def normalize(
+            default_level: LOG_LEVEL,
+            default_msg: str,
+            option: Optional[Tuple[LOG_LEVEL, str]] | str | bool | None,
+        ) -> Tuple[LOG_LEVEL, str] | None:
+        """
+        Normalize the options to specified args and make the input to `Tuple[LOG_LEVEL, str] | None`.
+        Returns None on negative inputs.
+
+        Args:
+            default_level(LOG_LEVEL): default level on str, bool inputs
+            default_msg(str): default msg on bool inputs
+            option(Optional[Tuple[LOG_LEVEL, str]] | str | bool | None): The option to normalize
+        
+        Returns:
+            (Tuple[LOG_LEVEL, str] | None): Normalized output for logging wraper
+        """
+        if isinstance(option, bool) and option is True:
+            return (default_level, default_msg)
+
+        elif isinstance(option, str):
+            return (default_level, option)
 
-    if isinstance(after, str):
-        after = ('DEFAULT', after)
+        elif isinstance(option, tuple):
+            return option
 
-    def decorator(func):
+        elif not option or option is None:
+            return None
+
+    before = normalize('DEBUG', 'Calling {func} - args:{args}, kwargs:{kwargs}', before)
+    after = normalize('INFO', 'Function {func} ended', after)
+    on_exception = normalize('ERROR', 'Error on {func}: {e}', on_exception)
+
+    def decorator(func: Callable) -> Callable:
         @wraps(func)
-        def wrapper(*args, **kwargs):
+        def wrapper(*args, **kwargs) -> Any:
             logger = logging.getLogger(func.__module__)
+            func_name = func.__name__
 
-            # Log before
-            if before:
-                level, message = before
-                logger.log(getattr(logging, level, logging.DEBUG), message)
+            fmt_context = {
+                'func': func_name,
+                'args': args,
+                'kwargs': kwargs,
+            }
 
-            # Log method parameters if requested
-            if show_args:
-                level = getattr(logging, show_args, logging.DEBUG)
-                logger.log(level, f"Arguments: args={args}, kwargs={kwargs}")
+            if before:
+                level, msg = before
+                logger.log(getattr(logging, level, logging.DEBUG), msg.format(**fmt_context))
 
-            result = func(*args, **kwargs)
+            try:
+                result = func(*args, **kwargs)
+            except Exception as e:
+                if on_exception:
+                    level, msg = on_exception
+                    fmt_context['e'] = e
+                    logger.log(getattr(logging, level, logging.ERROR), msg.format(**fmt_context))
+                raise e
 
-            # Log after
             if after:
-                level, message = after
-                logger.log(getattr(logging, level, logging.INFO), message)
+                level, msg = after
+                logger.log(getattr(logging, level, logging.INFO), msg.format(**fmt_context))
 
             return result
         return wrapper

test/test_logwrap.py

@@ -9,7 +9,7 @@ def my_new_func():
 
 def test_with_custom_leveling():
     # Example with Custom Levels
-    @logwrap(before=('INFO', 'Function starting'), after=('INFO', 'Function ended'), show_args='DEBUG')
+    @logwrap(before=('INFO', 'Function starting'), after=('INFO', 'Function ended'))
     def my_func(my_arg, my_kwarg=None):
         ...
     my_func('hello', my_kwarg=123) # calling the function