Python wrapt.decorator() Examples

def removed_kwarg(old_name, message=None,
                  version=None, removal_version=None, stacklevel=3,
                  category=None):
    """Decorates a kwarg accepting function to deprecate a removed kwarg."""

    prefix = "Using the '%s' argument is deprecated" % old_name
    out_message = _utils.generate_message(
        prefix, postfix=None, message=message, version=version,
        removal_version=removal_version)

    @wrapt.decorator
    def wrapper(f, instance, args, kwargs):
        if old_name in kwargs:
            _utils.deprecation(out_message,
                               stacklevel=stacklevel, category=category)
        return f(*args, **kwargs)

    return wrapper 

def require(*permissions):
        """
        Wrap a view function, verifying that the user hsa all of the specified
        permissions.
        """
        assert permissions, "Must specify at least one permission"
        for perm in permissions:
            if not perm.exists():
                raise RuntimeError(
                    "Cannot require undocumented permission %s" % (perm,))

        @wrapt.decorator
        def req(wrapped, instance, args, kwargs):
            if not can(*permissions):
                # redirect browsers when the user is not logged in, but
                # just return 403 to REST clients
                if util.is_browser() and current_user.is_anonymous:
                    return current_app.login_manager.unauthorized()
                else:
                    abort(403)
            return wrapped(*args, **kwargs)
        return req 

def verify_status(valid_status_list, raise_=False):
    """
    Decorator to help with debugging of state transitions
    If a decorated is called when the actors status is not in valid_status_list
    it will log (or raise exception if raise_ is True) the attempt.
    """
    @wrapt.decorator
    def wrapper(wrapped, instance, args, kwargs):
        # Exclude the instance variables added by superclasses
        if not instance.fsm.disable_state_checks and instance.fsm.state() not in valid_status_list:
            msg = "Invalid status %s for operation %s" % (instance.fsm, wrapped.__name__)
            if raise_:
                raise Exception(msg)
            else:
                _log.info(msg)
        x = wrapped(*args, **kwargs)
        return x
    return wrapper 

def versionadded(reason="", version=""):
    """
    This decorator can be used to insert a "versionadded" directive
    in your function/class docstring in order to documents the
    version of the project which adds this new functionality in your library.

    :param str reason:
        Reason message which documents the addition in your library (can be omitted).

    :param str version:
        Version of your project which adds this feature.
        If you follow the `Semantic Versioning <https://semver.org/>`_,
        the version number has the format "MAJOR.MINOR.PATCH", and,
        in the case of a new functionality, the "PATCH" component should be "0".

    :return: the decorated function.
    """
    adapter = SphinxAdapter('versionadded', reason=reason, version=version)

    # noinspection PyUnusedLocal
    @wrapt.decorator(adapter=adapter)
    def wrapper(wrapped, instance, args, kwargs):
        return wrapped(*args, **kwargs)

    return wrapper 

def versionchanged(reason="", version=""):
    """
    This decorator can be used to insert a "versionchanged" directive
    in your function/class docstring in order to documents the
    version of the project which modifies this functionality in your library.

    :param str reason:
        Reason message which documents the modification in your library (can be omitted).

    :param str version:
        Version of your project which modifies this feature.
        If you follow the `Semantic Versioning <https://semver.org/>`_,
        the version number has the format "MAJOR.MINOR.PATCH".

    :return: the decorated function.
    """
    adapter = SphinxAdapter('versionchanged', reason=reason, version=version)

    # noinspection PyUnusedLocal
    @wrapt.decorator(adapter=adapter)
    def wrapper(wrapped, instance, args, kwargs):
        return wrapped(*args, **kwargs)

    return wrapper 

def on(f):
    """Decorator for action handlers.

    The action name is inferred from the function name.

    This also decorates the method with `tornado.gen.coroutine` so that
    `~tornado.concurrent.Future` can be yielded.
    """
    action = f.__name__
    f.action = action

    @wrapt.decorator
    @tornado.gen.coroutine
    def _execute(wrapped, instance, args, kwargs):
        return wrapped(*args, **kwargs)

    return _execute(f) 

def trace_xray_subsegment(skip_args=False):
    """Can be applied to any function or method to be traced by X-Ray.

    If `skip_args` is True, the arguments of the function won't be sent to
    X-Ray.
    """

    @wrapt.decorator
    def wrapper(wrapped, instance, args, kwargs):
        metadata_extractor = (
            noop_function_metadata if skip_args else extract_function_metadata
        )
        return generic_xray_wrapper(
            wrapped,
            instance,
            args,
            kwargs,
            name=get_function_name,
            namespace="local",
            metadata_extractor=metadata_extractor,
        )

    return wrapper 

def reconfirm_auth(func, *args, **kwargs):
    '''
    A decorator to require the user to reconfirm their login.  Useful for sensitive pages.
    '''
    from allura.lib.plugin import AuthenticationProvider

    if request.POST.get('password'):
        if AuthenticationProvider.get(request).validate_password(c.user, request.POST['password']):
            session['auth-reconfirmed'] = datetime.utcnow()
            session.save()
            kwargs.pop('password', None)
        else:
            c.form_errors['password'] = 'Invalid password.'

    allowed_timedelta = timedelta(seconds=asint(config.get('auth.reconfirm.seconds', 60)))
    last_reconfirm = session.get('auth-reconfirmed', datetime.min)
    if datetime.utcnow() - last_reconfirm <= allowed_timedelta:
        return func(*args, **kwargs)
    else:
        return render({}, 'jinja', "allura:templates/reconfirm_auth.html") 

def _instance_method_wrapper(wrapped, rep_invariants=None):
    if rep_invariants is None:
        return check_contracts

    @wrapt.decorator
    def wrapper(wrapped, instance, args, kwargs):
        init = getattr(instance, '__init__')
        try:
            r = _check_function_contracts(wrapped, instance, args, kwargs)
            _check_invariants(instance, rep_invariants, init.__globals__)
            _check_class_type_annotations(instance)
        except AssertionError as e:
            raise AssertionError(str(e)) from None
        else:
            return r

    return wrapper(wrapped) 

def tracked_requests():
    """
    Gather all TrackedRequests that are buffered during a test into a list.
    """
    requests = []

    @wrapt.decorator
    def capture_requests(wrapped, instance, args, kwargs):
        if instance.is_real_request and not instance.is_ignored():
            requests.append(instance)
        return wrapped(*args, **kwargs)

    orig = TrackedRequest.finish
    TrackedRequest.finish = capture_requests(orig)
    try:
        yield requests
    finally:
        TrackedRequest.finish = orig 

def install_background_instrumentation():
    global background_instrumentation_installed
    if background_instrumentation_installed:
        return
    background_instrumentation_installed = True

    @wrapt.decorator
    async def wrapped_background_call(wrapped, instance, args, kwargs):
        tracked_request = TrackedRequest.instance()
        tracked_request.is_real_request = True
        tracked_request.start_span(
            operation="Job/{}.{}".format(
                instance.func.__module__, instance.func.__qualname__
            )
        )
        try:
            return await wrapped(*args, **kwargs)
        finally:
            tracked_request.stop_span()

    BackgroundTask.__call__ = wrapped_background_call(BackgroundTask.__call__) 

def renamed_kwarg(old_name, new_name, message=None,
                  version=None, removal_version=None, stacklevel=3,
                  category=None, replace=False):
    """Decorates a kwarg accepting function to deprecate a renamed kwarg."""

    prefix = _KWARG_RENAMED_PREFIX_TPL % old_name
    postfix = _KWARG_RENAMED_POSTFIX_TPL % new_name
    out_message = _utils.generate_message(
        prefix, postfix=postfix, message=message, version=version,
        removal_version=removal_version)

    @wrapt.decorator
    def decorator(wrapped, instance, args, kwargs):
        if old_name in kwargs:
            _utils.deprecation(out_message,
                               stacklevel=stacklevel, category=category)
            if replace:
                kwargs.setdefault(new_name, kwargs.pop(old_name))
        return wrapped(*args, **kwargs)

    return decorator 

def wrapped(signal_name, sender=dispatcher.Anonymous, safe=False):
    """A decorator for wrapping function in signal dispatch."""
    @wrapt.decorator
    def signal_wrapped(wrapped_func, _, args, kwargs):
        def signal_wrapper(*args, **kwargs):
            with wrap(signal_name, sender, safe):
                return wrapped_func(*args, **kwargs)

        return signal_wrapper(*args, **kwargs)

    return signal_wrapped 

def retry_zotero(wrapped, _instance, args, kwargs):
    """
    Retry the wrapped function if the Zotero API call fails.

    Caution: This decorator should only be used on simple functions, as the
    whole function is called repeatedly as long a Zotero fails.
    """
    attempts = 1
    while True:
        try:
            return wrapped(*args, **kwargs)
        except (
                requests.exceptions.ConnectionError,
                zotero_errors.HTTPError,
                zotero_errors.UnsupportedParams
        ) as e:
            current_app.logger.exception(e)
            if attempts < current_app.config['KERKO_ZOTERO_MAX_ATTEMPTS']:
                current_app.logger.warning(
                    "The Zotero API call has failed in {func}. "
                    "New attempt in {wait} seconds...".format(
                        func=wrapped.__name__,
                        wait=current_app.config['KERKO_ZOTERO_WAIT']
                    )
                )
                attempts += 1
                sleep(current_app.config['KERKO_ZOTERO_WAIT'])
            else:
                current_app.logger.error(
                    "The maximum number of API call attempts to Zotero has "
                    "been reached. Stopping."
                )
                raise 

def _permute_args(index_permutation):
    @wrapt.decorator
    def wrapper(wrapped, _, args, kwargs):
        if kwargs:
            raise _ForbiddenKwargs
        if len(index_permutation) != len(args):
            raise _WrongNumberOfArgs
        return wrapped(*_permute(args, index_permutation))

    return wrapper 

def manage(include=None, exclude=None):

    """
    Decorator for Actor::init() providing automatic management of state variables.
    Usage:
        @manage()                     # Manage every instance variable known upon completion of __init__
        @manage(include = [])         # Manage nothing
        @manage(include = [foo, bar]) # Manage self.foo and self.bar only. Equivalent to @manage([foo, bar])
        @manage(exclude = [foo, bar]) # Manage everything except self.foo and self.bar
        @manage(exclude = [])         # Same as @manage()
        @manage(<list>)               # Same as @manage(include = <list>)

    N.B. If include and exclude are both present, exclude will be disregarded.

    """

    if include and type(include) is not list or exclude and type(exclude) is not list:
        raise Exception("@manage decorator: Must use list as argument")

    include_set = set(include) if include else set()
    exclude_set = set(exclude) if exclude else set()

    # Using wrapt since we need to preserve the signature of the wrapped signature.
    # See http://wrapt.readthedocs.org/en/latest/index.html
    # FIXME: Since we use wrapt here, we might as well use it in guard and condition too.
    @wrapt.decorator
    def wrapper(wrapped, instance, args, kwargs):
        # Exclude the instance variables added by superclasses
        exclude_set.update(instance.__dict__)
        x = wrapped(*args, **kwargs)
        if include is None:
            # include set not given, so construct the implicit include set
            include_set.update(instance.__dict__)
            include_set.remove('_managed')
            include_set.difference_update(exclude_set)
        instance._managed.update(include_set)
        return x
    return wrapper 

def synchronized(lock):
    @wrapt.decorator
    def wrap(wrapper, instance, args, kwargs):
        with lock:
            return wrapper(*args, **kwargs)
    return wrap 

def wait_container_is_ready():
    """
    Wait until container is ready.
    Function that spawn container should be decorated by this method
    Max wait is configured by config. Default is 120 sec.
    Polling interval is 1 sec.
    :return:
    """

    @wrapt.decorator
    def wrapper(wrapped, instance, args, kwargs):
        exception = None
        print(crayons.yellow("Waiting to be ready..."))
        with blindspin.spinner():
            for _ in range(0, config.MAX_TRIES):
                try:
                    return wrapped(*args, **kwargs)
                except Exception as e:
                    time.sleep(config.SLEEP_TIME)
                    exception = e
            raise TimeoutException(
                """Wait time exceeded {0} sec.
                    Method {1}, args {2} , kwargs {3}.
                     Exception {4}""".format(config.MAX_TRIES,
                                             wrapped.__name__,
                                             args, kwargs, exception))

    return wrapper 

def add_port_docs(wrapped=None, instance=None, value=''):
    if wrapped is None:
        return functools.partial(add_port_docs, value=value)

    @wrapt.decorator
    def wrapper(wrapped, instance=None, args=None, kwargs=None):
        return wrapped(*args, **kwargs)

    decorated = wrapper(wrapped)
    try:
        port_2_ntype = decorated(instance)
    except:
        port_2_ntype = None

    port_description = ""
    if port_2_ntype is not None:
        for port, ntype in port_2_ntype.items():
            port_description += "* *" + port + "* : " + str(ntype)
            port_description += "\n\n"

    __doc__ = _normalize_docstring(wrapped.__doc__) + '\n\n' + str(port_description)
    __doc__ = _normalize_docstring(__doc__)

    wrapt.FunctionWrapper.__setattr__(decorated, "__doc__", __doc__)

    return decorated 

def required_parameters(*parameters):
    """A decorator for views with required parameters.

    Returns a 400 if parameters are not provided by the client.

    Warning: when applied, it turns the request object into a data one,
    as first parameter of the returned function to avoid parsing the
    JSON data twice.
    """
    @wrapt.decorator
    def wrapper(wrapped, instance, args, kwargs):
        args = list(args)
        request = args[0]
        try:
            data = data_from_request(request)
        except ValueError as error:
            return 400, 'Incorrect parameters: {error}'.format(error=error)
        data.update(flatten_get_parameters(request))
        for parameter in parameters:
            if parameter not in data:
                log(('"{parameter}" parameter not found in {data}'
                     .format(data=data, parameter=parameter)))
                return 400, ('Please specify a "{parameter}" parameter.'
                             .format(data=data, parameter=parameter))
        args[0] = data
        return wrapped(*args, **kwargs)
    return wrapper 

def __init__(self, states, initial, transitions, hooks=None, allow_invalid_transitions=True,
                     disable_transition_checks=False, disable_state_checks=False):
            self.states = states
            self._state = initial
            self.transitions = transitions
            self.hooks = hooks or {}
            self.allow_invalid_transitions = allow_invalid_transitions
            self.disable_transition_checks = disable_transition_checks
            # disable_state_checks is used in the verify_status decorator
            self.disable_state_checks = disable_state_checks 

def base_hook(sender_callable, hash_function, **dkwargs):

    @wrapt.decorator
    def wrapper(wrapped, instance, args, kwargs):
        """ This should calls the sender_callable.
            kwargs needs to include an 'creator' key
        """

        # If sender_callable isn't valid, stop early for easy debugging
        if not callable(sender_callable):
            raise SenderNotCallable(sender_callable)

        # Call the hash function and save result to a hash_value argument
        hash_value = None
        if hash_function is not None:
            hash_value = hash_function()

        ##################################
        # :wrapped: hooked function delivering a payload
        # :dkwargs: name of the event being called
        # :hash_value: hash_value to determine the uniqueness of the payload
        # :args: Argument list for the wrapped function
        # :kwargs: Keyword arguments for the wrapped function. Must include 'creator'

        # Send the hooked function
        status = sender_callable(wrapped, dkwargs, hash_value, *args, **kwargs)

        # Status can be anything:
        #   * The result of a synchronous sender
        #   * A generic status message for asynchronous senders
        #   * Returns the response of the hash_function
        return status

    return wrapper


# This is the hook everyone wants to use 

def __doc__(self):
        doc = getattr(self.wrapped, "__doc__", None)
        return "<wrapped by the cachedproperty decorator>%s" % (
            "\n%s" % doc if doc else ""
        ) 

def cached(func, instance, args, kwargs):
    """Simple decorator to cache result of method calls without args."""
    cache = getattr(instance, "__cache", None)
    if cache is None:
        instance.__cache = cache = {}
    try:
        return cache[func]
    except KeyError:
        cache[func] = result = func(*args, **kwargs)
        return result 

def _inference_tip_cached(func, instance, args, kwargs, _cache={}):
    """Cache decorator used for inference tips"""
    node = args[0]
    try:
        return iter(_cache[func, node])
    except KeyError:
        result = func(*args, **kwargs)
        # Need to keep an iterator around
        original, copy = itertools.tee(result)
        _cache[func, node] = list(copy)
        return original


# pylint: enable=dangerous-default-value 

def retry(total, backoff_factor, backoff_max, exceptions):

    total = max(total, 1)
    backoff_factor = max(backoff_factor, 0)
    backoff_max = float('inf') if backoff_max is None else max(0, backoff_max)

    @wrapt.decorator
    def wrapper(wrapped, instance, args, kwargs):
        errors = 0
        while True:
            try:
                return wrapped(*args, **kwargs)
            except exceptions:
                errors += 1

                if errors > total:
                    raise

                if errors >= 2:
                    backoff_value = backoff_factor * (2 ** (errors - 1))
                else:
                    backoff_value = 0
                backoff_value = min(backoff_value, backoff_max)

                sleep(backoff_value)

    return wrapper 

def transaction_retry(wrapped=None, session=None, total=1,
                      backoff_factor=0, backoff_max=None):

    if wrapped is None:
        return functools.partial(
            transaction_retry, session=session,
            total=total,
            backoff_factor=backoff_factor,
            backoff_max=backoff_max)

    @wrapt.decorator
    def wrapper(wrapped, instance, args, kwargs):

        @retry(total, backoff_factor, backoff_max, exc.OperationalError)
        def run_or_rollback():
            try:
                return wrapped(*args, **kwargs)
            except exc.OperationalError as exception:
                if exception.connection_invalidated:
                    if isinstance(session, operator.attrgetter):
                        session(instance).rollback()
                    elif session:
                        session.rollback()
                raise

        return run_or_rollback()

    return wrapper(wrapped)  # pylint: disable=E1120 

def updated_kwarg_default_value(name, old_value, new_value, message=None,
                                version=None, stacklevel=3,
                                category=FutureWarning):

    """Decorates a kwarg accepting function to change the default value"""

    prefix = _KWARG_UPDATED_PREFIX_TPL % (name, new_value)
    postfix = _KWARG_UPDATED_POSTFIX_TPL % old_value
    out_message = _utils.generate_message(
        prefix, postfix=postfix, message=message, version=version)

    def decorator(f):
        sig = signature(f)
        varnames = list(sig.parameters.keys())

        @wrapt.decorator
        def wrapper(wrapped, instance, args, kwargs):
            explicit_params = set(
                varnames[:len(args)] + list(kwargs.keys())
            )
            allparams = set(varnames)
            default_params = set(allparams - explicit_params)
            if name in default_params:
                _utils.deprecation(out_message,
                                   stacklevel=stacklevel, category=category)
            return wrapped(*args, **kwargs)

        return wrapper(f)

    return decorator 

def moved_class(new_class, old_class_name, old_module_name,
                message=None, version=None, removal_version=None,
                stacklevel=3, category=None):
    """Deprecates a class that was moved to another location.

    This creates a 'new-old' type that can be used for a
    deprecation period that can be inherited from. This will emit warnings
    when the old locations class is initialized, telling where the new and
    improved location for the old class now is.
    """

    if not inspect.isclass(new_class):
        _qual, type_name = _utils.get_qualified_name(type(new_class))
        raise TypeError("Unexpected class type '%s' (expected"
                        " class type only)" % type_name)

    old_name = ".".join((old_module_name, old_class_name))
    new_name = _utils.get_class_name(new_class)
    prefix = _CLASS_MOVED_PREFIX_TPL % (old_name, new_name)
    out_message = _utils.generate_message(
        prefix, message=message, version=version,
        removal_version=removal_version)

    def decorator(f):

        @six.wraps(f, assigned=_utils.get_assigned(f))
        def wrapper(self, *args, **kwargs):
            _utils.deprecation(out_message, stacklevel=stacklevel,
                               category=category)
            return f(self, *args, **kwargs)

        return wrapper

    old_class = type(old_class_name, (new_class,), {})
    old_class.__module__ = old_module_name
    old_class.__init__ = decorator(old_class.__init__)
    return old_class 

def subsegment_decorator(wrapped, instance, args, kwargs):
    decorated_func = wrapt.decorator(wrapped)(*args, **kwargs)
    set_as_recording(decorated_func, wrapped)
    return decorated_func