Environments with fabric

cmd

$ fab -l
Available commands:

    production.cleannn
    production.deploy
    staging.deploy

fabfile.py

from task import task, environment
from fabric.api import env
from fabric.operations import run

@environment
def production():
    env.host_string = 'test@production.org'

@environment
def staging():
    env.host_string = 'test@staging.org'

@task()
def deploy():
    run('whoami')

@task(name='cleannn', environments=['production'])
def clean():
    run('echo \o/')

task.py

import types

from fabric import tasks
from fabric import state
from fabric.task_utils import _Dict
from functools import wraps
from fabric.api import env

environments = {}
state.env.new_style_tasks = True

def environment(config):
    """
    Specify a new environment
    """
    @wraps(config)
    def inner():
        env.environment = config.__name__
        return config()

    environments[config.__name__] = inner

    return inner

def task(*args, **kwargs):
    """
    Specify a new task
    """
    invoked = bool(not args or kwargs)
    task_class = kwargs.pop("task_class", tasks.WrappedCallableTask)

    if not invoked:
        func, args = args[0], ()

    def wrapper(func):
        task_name  = kwargs.pop("name", func.__name__)
        env_list = kwargs.pop("environments", [])

        for name, env in environments.iteritems():
            if env_list and name not in env_list:
                continue

            env_task_name = '%s.%s' % (name, task_name)
            function = wrapenv(env, func, env_task_name)
            insert_command(env_task_name, task_class(function, *args, **kwargs), state.commands)

        return func

    return wrapper if invoked else wrapper(func)

def wrapenv(envconfig, function, task_name):
    """
    Wrapper for the task to call environment config in first place then the command
    """
    @wraps(function)
    def envfunc(*args, **kwargs):
        envconfig()
        env.current_task_name = task_name
        function(*args, **kwargs)

    return envfunc

def insert_command(name, command, mapping):
    """
    Insert a command into fabric
    """
    key, _, rest = name.partition('.')

    if not rest:
        mapping[key] = command
        return

    if key not in mapping:
        mapping[key] = _Dict()

    return insert_command(rest, command, mapping[key])