def after_func_accept_retry_state(fn):
    """Wrap "after" function to accept "retry_state"."""
    if not six.callable(fn):
        return fn

    if func_takes_retry_state(fn):
        return fn

    @_utils.wraps(fn)
    def wrapped_after_sleep_func(retry_state):
        # func, trial_number, trial_time_taken
        warn_about_non_retry_state_deprecation('after', fn, stacklevel=4)
        return fn(
            retry_state.fn,
            retry_state.attempt_number,
            retry_state.seconds_since_start)
    return wrapped_after_sleep_func
Wrap "after" function to accept "retry_state".