deb-python-taskflow/taskflow/decorators.py

# -*- coding: utf-8 -*-

# vim: tabstop=4 shiftwidth=4 softtabstop=4

#    Copyright (C) 2012 Yahoo! Inc. All Rights Reserved.
#
#    Licensed under the Apache License, Version 2.0 (the "License"); you may
#    not use this file except in compliance with the License. You may obtain
#    a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
#    License for the specific language governing permissions and limitations
#    under the License.

import collections
import functools
import inspect

from taskflow import utils

# These arguments are ones that we will skip when parsing for requirements
# for a function to operate (when used as a task).
AUTO_ARGS = ('self', 'context', 'cls')


def _take_arg(a):
    if a in AUTO_ARGS:
        return False
    # In certain decorator cases it seems like we get the function to be
    # decorated as an argument, we don't want to take that as a real argument.
    if isinstance(a, collections.Callable):
        return False
    return True


def wraps(fn):
    """This will not be needed in python 3.2 or greater which already has this
    built-in to its functools.wraps method."""

    def wrapper(f):
        f = functools.wraps(fn)(f)
        f.__wrapped__ = getattr(fn, '__wrapped__', fn)
        return f

    return wrapper


def task(*args, **kwargs):
    """Decorates a given function and ensures that all needed attributes of
    that function are set so that the function can be used as a task."""

    def decorator(f):

        def noop(*args, **kwargs):
            pass

        f.revert = kwargs.pop('revert_with', noop)

        # Sets the version of the task.
        version = kwargs.pop('version', (1, 0))
        f = versionize(*version)(f)

        # Attach any requirements this function needs for running.
        requires_what = kwargs.pop('requires', [])
        f = requires(*requires_what, **kwargs)(f)

        # Attach any optional requirements this function needs for running.
        optional_what = kwargs.pop('optional', [])
        f = optional(*optional_what, **kwargs)(f)

        # Attach any items this function provides as output
        provides_what = kwargs.pop('provides', [])
        f = provides(*provides_what, **kwargs)(f)

        # Associate a name of this task that is the module + function name.
        f.name = "%s.%s" % (f.__module__, f.__name__)

        @wraps(f)
        def wrapper(*args, **kwargs):
            return f(*args, **kwargs)

        return wrapper

    # This is needed to handle when the decorator has args or the decorator
    # doesn't have args, python is rather weird here...
    if kwargs or not args:
        return decorator
    else:
        if isinstance(args[0], collections.Callable):
            return decorator(args[0])
        else:
            return decorator


def versionize(major, minor=None):
    """A decorator that marks the wrapped function with a major & minor version
    number."""

    if minor is None:
        minor = 0

    def decorator(f):
        f.__version__ = (major, minor)

        @wraps(f)
        def wrapper(*args, **kwargs):
            return f(*args, **kwargs)

        return wrapper

    return decorator


def optional(*args, **kwargs):
    """Attaches a set of items that the decorated function would like as input
    to the functions underlying dictionary."""

    def decorator(f):
        if not hasattr(f, 'optional'):
            f.optional = set()

        f.optional.update([a for a in args if _take_arg(a)])

        @wraps(f)
        def wrapper(*args, **kwargs):
            return f(*args, **kwargs)

        return wrapper

    # This is needed to handle when the decorator has args or the decorator
    # doesn't have args, python is rather weird here...
    if kwargs or not args:
        return decorator
    else:
        if isinstance(args[0], collections.Callable):
            return decorator(args[0])
        else:
            return decorator


def requires(*args, **kwargs):
    """Attaches a set of items that the decorated function requires as input
    to the functions underlying dictionary."""

    def decorator(f):
        if not hasattr(f, 'requires'):
            f.requires = set()

        if kwargs.pop('auto_extract', True):
            inspect_what = getattr(f, '__wrapped__', None)

            if not inspect_what:
                inspect_what = utils.get_wrapped_function(f)

            f_args = inspect.getargspec(inspect_what).args
            f.requires.update([a for a in f_args if _take_arg(a)])

        f.requires.update([a for a in args if _take_arg(a)])

        @wraps(f)
        def wrapper(*args, **kwargs):
            return f(*args, **kwargs)

        return wrapper

    # This is needed to handle when the decorator has args or the decorator
    # doesn't have args, python is rather weird here...
    if kwargs or not args:
        return decorator
    else:
        if isinstance(args[0], collections.Callable):
            return decorator(args[0])
        else:
            return decorator


def provides(*args, **kwargs):
    """Attaches a set of items that the decorated function provides as output
    to the functions underlying dictionary."""

    def decorator(f):
        if not hasattr(f, 'provides'):
            f.provides = set()

        f.provides.update([a for a in args if _take_arg(a)])

        @wraps(f)
        def wrapper(*args, **kwargs):
            return f(*args, **kwargs)

        return wrapper

    # This is needed to handle when the decorator has args or the decorator
    # doesn't have args, python is rather weird here...
    if kwargs or not args:
        return decorator
    else:
        if isinstance(args[0], collections.Callable):
            return decorator(args[0])
        else:
            return decorator