from aws_xray_sdk.core import xray_recorder as xray


def xray_capture(func):
    """Wraps xray.capture(), allowing a class to opt in/out"""

    def wrapper(*args, **kwargs):
        if (
            len(args) > 0
            and isinstance(args[0], object)
            and callable(getattr(args[0], func.__name__, None))
            and not getattr(args[0], "_xray_capture", False)
        ):
            # "func" is most likely an object function on an object
            # that has not set its "_xray_capture" to True
            # Do not allow xray.capture()
            return func(*args, **kwargs)
        else:
            return xray.capture()(func)(*args, **kwargs)

    return wrapper
