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
import types

# 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 is_decorated(functor):
    if not isinstance(functor, (types.MethodType, types.FunctionType)):
        return False
    return getattr(extract(functor), '__task__', False)


def extract(functor):
    # Extract the underlying functor if its a method since we can not set
    # attributes on instance methods, this is supposedly fixed in python 3
    # and later.
    #
    # TODO(harlowja): add link to this fix.
    assert isinstance(functor, (types.MethodType, types.FunctionType))
    if isinstance(functor, types.MethodType):
        return functor.__func__
    else:
        return functor


def _mark_as_task(functor):
    setattr(functor, '__task__', True)


def _get_wrapped(function):
    """Get the method at the bottom of a stack of decorators."""

    if hasattr(function, '__wrapped__'):
        return getattr(function, '__wrapped__')

    if not hasattr(function, 'func_closure') or not function.func_closure:
        return function

    def _get_wrapped_function(function):
        if not hasattr(function, 'func_closure') or not function.func_closure:
            return None

        for closure in function.func_closure:
            func = closure.cell_contents

            deeper_func = _get_wrapped_function(func)
            if deeper_func:
                return deeper_func
            elif hasattr(closure.cell_contents, '__call__'):
                return closure.cell_contents

    return _get_wrapped_function(function)


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 locked(f):

    @wraps(f)
    def wrapper(self, *args, **kwargs):
        with self._lock:
            return f(self, *args, **kwargs)

    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):
        w_f = extract(f)

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

        # Mark as being a task
        _mark_as_task(w_f)

        # By default don't revert this.
        w_f.revert = kwargs.pop('revert_with', noop)

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

        # 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)

        @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):
        w_f = extract(f)
        w_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):
        w_f = extract(f)

        if not hasattr(w_f, 'optional'):
            w_f.optional = set()

        w_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):
        w_f = extract(f)

        if not hasattr(w_f, 'requires'):
            w_f.requires = set()

        if kwargs.pop('auto_extract', True):
            inspect_what = _get_wrapped(f)
            f_args = inspect.getargspec(inspect_what).args
            w_f.requires.update([a for a in f_args if _take_arg(a)])

        w_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):
        w_f = extract(f)

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

        w_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