Fabric decorators

fab_decorators.py

import fabric
from fabric.api import env
from functools import wraps

###############################################################################
# Decorators
###############################################################################


def limit_roles(*role_list):
    """
    Decorator defining a list of role names that the task is limited to.

    If run with a host outside the role_list - it will just return and not
    execute the function.

    A role is simply defined as a key in `env` whose value is a list of one or
    more host connection strings. For example, the following will ensure that,
    ``my_func`` will only be executed for the hosts listed in either the 
    ``www`` or ``test`` roles:: 


        env.roledefs.update({
            'www': ['www1', 'www2'],
            'db': ['db1']
            'test': ['test1'],
        })

        @limit_roles('webserver', 'dbserver')
        def my_func():
            pass

        .. note::
            This differs from the @roles decorator in that it uses the current
            role to inspect the hosts, rather than assigning the role.  So
            running ``my_func`` under the ``db`` role will just return.
    """
    def attach_roles(func):
        @wraps(func)
        def inner_decorator(*args, **kwargs):
            if env.host_string not in inner_decorator._available_hosts:
                return
            return func(*args, **kwargs)
        _role_list = role_list
        if isinstance(_role_list, basestring):
            _role_list = [_role_list]
        hosts = []
        for role in _role_list:
            hosts.extend(env.roledefs.get(role, []))
        inner_decorator._available_hosts = hosts
        return inner_decorator
    return attach_roles


def memoize_question(question, default=False):
    """
    Decorator to memoize the answer of a question, so you don't get asked for 
    each host in a list. 

    For example the following example asks if you want to restart memcached::

    from fabric.colors import green

    @task
    @memoize_question(green("Restart Memcached?"), default=True)
    def restart_memcached():
        sudo('service memcached restart')


    """
    def attach_questions(func):
        @wraps(func)
        def inner_decorator(*args, **kwargs):
            name = func.__name__
            asked_attr = "_asked_%s" % name
            answer_attr = "_answer_%s" % name

            if asked_attr not in env:
                setattr(env, asked_attr, True)
                setattr(env, answer_attr, False)

                if fabric.contrib.console.confirm(question, default=default):
                    setattr(env, answer_attr, True)

            if getattr(env, answer_attr, False):
                return func(*args, **kwargs)
            return None
        return inner_decorator
    return attach_questions