Ryan

import types
from collections import defaultdict
import inspect

class Contract(object):
    """Class to contain preconditions and postconditions, as well as expose
    methods as decorators."""
    def __init__(self, func=None):
        self.do(func)
        self.invariants = []
        self._invariant_called = False
        self.preconditions = []
        self.postconditions = []

    def do(self, func):
        """Decorator to associate a method with a contract."""
        self.__doc__ = func.__doc__
        self.method = func
        return self

    def ensure(self, func):
        """Decorator to set post conditions."""
        self.postconditions.append(func)
        return self

    def require(self, func):
        """Decorator to set pre conditions."""
        self.preconditions.append(func)
        return self

    def before(self, *args, **kwargs):
        """Precondition and invariants checked before method invocation."""
        precondition_errors = []
        for precondition in  self.preconditions:
            # 'OR' the preconditions together. This is allowed to pass as
            # long as one of them passes.
            try:
                precondition(*args, **kwargs)
            except AssertionError, err:
                precondition_errors.append(err)
            if len(precondition_errors) == len(self.preconditions):
                raise precondition_errors[0]

    def after(self, returned, *args, **kwargs):
        """Postcondition and invariants checked after method call."""
        # This is equivalent to 'AND'ing the postconditions together. If one
        # fails, the entire postcondition check fails. (This is
        # short-circuited).
        for postcondition in self.postconditions:
            postcondition(returned, *args, **kwargs)


    def __get__(self, obj, type=None):
        if obj is None:
            return self
        return types.MethodType(self, obj)

    def __call__(self, *args, **kwargs):
        """Wrap method in precondition, postcondition and invariant calls."""
        if self.preconditions:
            self.before(*args, **kwargs)
        if self.method:
            returned = self.method(*args, **kwargs)
        else:
            raise NotImplementedError("Method not created for contract")
        self.after(returned, *args, **kwargs)

        if self.invariants and not self._invariant_called:
            self._invariant_called = True
            for invariant in self.invariants:
                invariant(args[0])
            self._invariant_called = False

        return returned


def contract(func=None):
    """Decorator to mark a function as having pre and post conditions."""
    return Contract(func)

class DbcMeta(type):

    def __init__(cls, name, bases, attributes):
        super(DbcMeta, cls).__init__(name, bases, attributes)

        super_attrs = defaultdict(list)
        for base in bases:
            for k, v in inspect.getmembers(base):
                if isinstance(v, Contract):
                    super_attrs[k].append(v)
                if k == "invariant":
                    invariants.append(v)

        cls.convert_methods(attributes, super_attrs)

    def convert_methods(cls, attributes, super_attrs):
        """Replace functions in attributes with wrappers.

        The attributes is modified in place.
        """

        # find methods with contracts.
        contracted_methods = []
        invariants = []
        if hasattr(cls, 'invariant'):
            invariants.append(getattr(cls, "invariant"))
        for k, v in attributes.iteritems():
            if isinstance(v, Contract):
                contracted_methods.append(k)

        # Add inherited pre and post conditions
        # TODO: handle inherited invariants, multiple inheritance, etc.
        for m in contracted_methods:
            contract = getattr(cls,m)
            preconditions = [precond for cont in super_attrs[m]
                                         for precond in cont.preconditions]
            postconditions = [postcond for cont in super_attrs[m]
                                           for postcond in cont.postconditions]
            if invariants:
                contract.invariants = invariants
            for precond in preconditions:
                contract.require(precond)
            for postcond in postconditions:
                contract.ensure(postcond)

            setattr(cls, m, contract)