Decorator factory

decorator_factory.py

import functools
from typing import Callable, Dict, Any

# def lambda_handler_decorator(decorator: Callable):
#     """Decorator factory for decorating Lambda handlers."""

#     @functools.wraps(decorator)
#     def final_decorator(func: Callable = None, **kwargs):
#         if func is None:
#             return functools.partial(final_decorator, **kwargs)

#         @functools.wraps(func)
#         def wrapper(event: Dict, context: Dict):
#             try:
#                 print(f"ANNOTATION - BEGINS: {decorator.__qualname__}")
#                 f = decorator(func, event, context, **kwargs)
#                 print(f"ANNOTATION - END: {decorator.__qualname__}")
#                 return f
#             except Exception as err:
#                 print(f"CAUGHT exception in decorator {decorator.__qualname__}!")
#                 print(f"CAUGHT exception in final dec {final_decorator.__qualname__}!")
#                 raise err

#         return wrapper

#     return final_decorator

import functools
import logging
import os
from contextlib import contextmanager
from typing import Callable

logger = logging.getLogger(__name__)
logger.setLevel(os.getenv("LOG_LEVEL", "INFO"))


def lambda_handler_decorator(decorator: Callable = None, trace_execution=False):
    """Decorator factory for decorating Lambda handlers.

    You can use lambda_handler_decorator to create your own middlewares,
    where your function signature follows: fn(handler, event, context)

    You can also set your own key=value params: fn(handler, event, context, option=value)

    Example
    -------
    **Create a middleware no params**

        from aws_lambda_powertools.utils import lambda_handler_decorator

        @lambda_handler_decorator
        def log_response(handler, event, context):
            any_code_to_execute_before_lambda_handler()
            response = handler(event, context)
            any_code_to_execute_after_lambda_handler()
            print(f"Lambda handler response: {response}")

        @log_response
        def lambda_handler(event, context):
            return True

    **Create a middleware with params**

        from aws_lambda_powertools.utils import lambda_handler_decorator

        @lambda_handler_decorator
        def obfuscate_sensitive_data(handler, event, context, fields=None):
            # Obfuscate email before calling Lambda handler
            if fields:
                for field in fields:
                    field = event.get(field, "")
                    event[field] = obfuscate_pii(field)

            response = handler(event, context)
            print(f"Lambda handler response: {response}")

        @obfuscate_sensitive_data(fields=["email"])
        def lambda_handler(event, context):
            return True
    """
    if decorator is None:
        return functools.partial(lambda_handler_decorator, trace_execution=trace_execution)

    trace_execution = trace_execution or os.getenv("POWERTOOLS_TRACE_MIDDLEWARES", False)

    if decorator.__name__ != decorator.__qualname__: # identify non-bounded class methods
        class wrapper():
            def __init__(self, func):
                functools.update_wrapper(self, decorator)
                self.decorator = decorator
                self.handler = func

            def __call__(self, event, context):
                print(self.decorator.__qualname__)
                print(self.handler.__qualname__)
                return self.decorator(self, self.handler, event, context)

        return wrapper
    else:
        @functools.wraps(decorator)
        def final_decorator(func: Callable = None, *args, **kwargs):
            # If called with args return new func with args
            if func is None:
                return functools.partial(final_decorator, *args, **kwargs)

            @functools.wraps(func)
            def wrapper(event, context):
                try:
                    response = decorator(func, event, context, **kwargs)
                    return response
                except Exception as err:
                    logger.error(f"Caught exception in {decorator.__qualname__}")
                    raise err

            return wrapper

        return final_decorator


class Blah:    
    @lambda_handler_decorator
    def log_metric(self, handler: Callable, event: Dict, context: Any):
        print("Logging metric...")
        handler(event, context)
        print("Metric logged...")

@lambda_handler_decorator
def call_handler(handler, event, context):
    print("[CALL HANDLER] BEGIN")
    response = handler(event, context)
    print("[CALL HANDLER] END")
    return response


@lambda_handler_decorator
def also_call_handler(handler, event, context, log_response=None):
    print("[ALSO_CALL HANDLER] BEGIN")
    context.message = "call_handler"
    return handler(event, context)

decorator_factory_as_class.py

import functools
from typing import Callable, Dict, Any

import functools
import logging
import os
from contextlib import contextmanager
from typing import Callable

logger = logging.getLogger(__name__)
logger.setLevel(os.getenv("LOG_LEVEL", "INFO"))


def lambda_handler_decorator(decorator: Callable = None, trace_execution=False):
    """Decorator factory for decorating Lambda handlers.

    You can use lambda_handler_decorator to create your own middlewares,
    where your function signature follows: fn(handler, event, context)

    You can also set your own key=value params: fn(handler, event, context, option=value)

    Example
    -------
    **Create a middleware no params**

        from aws_lambda_powertools.utils import lambda_handler_decorator

        @lambda_handler_decorator
        def log_response(handler, event, context):
            any_code_to_execute_before_lambda_handler()
            response = handler(event, context)
            any_code_to_execute_after_lambda_handler()
            print(f"Lambda handler response: {response}")

        @log_response
        def lambda_handler(event, context):
            return True

    **Create a middleware with params**

        from aws_lambda_powertools.utils import lambda_handler_decorator

        @lambda_handler_decorator
        def obfuscate_sensitive_data(handler, event, context, fields=None):
            # Obfuscate email before calling Lambda handler
            if fields:
                for field in fields:
                    field = event.get(field, "")
                    event[field] = obfuscate_pii(field)

            response = handler(event, context)
            print(f"Lambda handler response: {response}")

        @obfuscate_sensitive_data(fields=["email"])
        def lambda_handler(event, context):
            return True
    """
    if decorator is None:
        return functools.partial(lambda_handler_decorator, trace_execution=trace_execution)

    trace_execution = trace_execution or os.getenv("POWERTOOLS_TRACE_MIDDLEWARES", False)

    class final_decorator():
        def __init__(self, func):
            functools.update_wrapper(self, decorator)
            self.decorator = decorator
            self.handler = func

        def __call__(self, event, context):
            print(self.decorator.__qualname__)
            print(self.handler.__qualname__)
            if decorator.__name__ != decorator.__qualname__:
                middleware = functools.partial(self.decorator, self, self.handler, event, context)
            else:
                middleware = functools.partial(self.decorator, self.handler, event, context)
            return middleware()
            # return self.decorator(self, self.handler, event, context)

    return final_decorator




class Blah:    
    @lambda_handler_decorator
    def log_metric(self, handler: Callable, event: Dict, context: Any):
        print("Logging metric...")
        handler(event, context)
        print("Metric logged...")

@lambda_handler_decorator
def call_handler(handler, event, context):
    print("[CALL HANDLER] BEGIN")
    response = handler(event, context)
    print("[CALL HANDLER] END")
    return response


@lambda_handler_decorator
def also_call_handler(handler, event, context, log_response=None):
    print("[ALSO_CALL HANDLER] BEGIN")
    context.message = "call_handler"
    return handler(event, context)