Decorator for Enforcing Argument Types

[Chris]

I'm not sure if this has been done before, but I couldn't easily find
any prior work on Google, so here I present a simple decorator for
documenting and verifying the type of function arguments.
Feedback/suggestions/criticism is welcome.

'''
2006.12.21 Created.
'''

import unittest
import inspect

def arguments(*args):
    '''A simple decorator for formally documenting
    and verifying argument types.

    usage:

        @arguments(type1, type2, [type3.1, type3.2, ..., type3.N],
type4, ..., typeN)
        def somefunc(arg1, arg2, arg3, arg4, ..., argN):
            do stuff
            return

    '''
    return lambda f:_Arguments(f, *args)

class _Arguments(object):
    # todo: extend to verify Zope interfaces
    def __init__(self, fn, *args):
        self.fn = fn

        # create argument type list
        self.arguments = []
        for arg in args:
            if not isinstance(arg, list):
                arg = list([arg])
            arg = set(arg)
            self.arguments.append(arg)

        # create name-to-index lookup
        argNames, varArgName, varkwName, defaults =
inspect.getargspec(fn)
        assert len(argNames) == len(self.arguments), 'list of argument
types must match the number of arguments'
        self.argNameToIndex = {}
        for i,name in enumerate(argNames):
            self.argNameToIndex[name] = i
            if defaults and i >= len(self.arguments)-len(defaults):
                # add default type to allowable types

self.arguments[i].add(type(defaults[i-(len(self.arguments)-len(defaults))]))

    def verify(self, value, i):
        '''Returns true if the value matches the allowable types
        for the ith argument.'''
        if not isinstance(i, int):
            if i not in self.argNameToIndex:
                raise Exception, 'unknown argument name: %s' % i
            i = self.argNameToIndex[i]
        return type(value) in self.arguments[i]

    def verifyAll(self, *values, **kvalues):
        '''Returns true if all values matche the allowable types
        for their corresponding arguments.'''
        for i,value in enumerate(values):
            if not self.verify(value, i):
                return False
        for name,value in kvalues.iteritems():
            if not self.verify(value, name):
                return False
        return True

    def __call__(self, *args, **kargs):
        assert self.verifyAll(*args, **kargs), 'argument types must be
in the form of %s' % self.arguments
        return self.fn(*args, **kargs)

class Test(unittest.TestCase):

    def test(self):

        @arguments(str, [int, float], list)
        def foo(abc, xyz, big=None):
            return '%s %s' % (abc, xyz)

        self.assertEqual(type(foo), _Arguments)
        self.assertEqual(len(foo.arguments), 3)
        self.assertEqual(foo.arguments[2], set([list, type(None)]))
        self.assertEqual(foo.verify('how', 0), True)
        self.assertEqual(foo.verify(123, 0), False)
        self.assertEqual(foo.verify(123, 1), True)
        self.assertEqual(foo.verify(1.23, 1), True)
        self.assertEqual(foo.verifyAll('how',123), True)
        self.assertEqual(foo.verifyAll(123,'how'), False)
        self.assertEqual(foo.verifyAll(abc='how',xyz=123), True)
        self.assertEqual(foo.verifyAll('how',xyz=123), True)
        self.assertEqual(foo.verifyAll('how',xyz='oeuuo'), False)
        self.assertEqual(foo.verifyAll('how',xyz=123,big=None), True)
        self.assertEqual(foo.verifyAll('how',xyz=123,big=[1,2,3]),
True)
        self.assertEqual(foo.verifyAll('how',123,[1,2,3]), True)
        self.assertEqual(foo.verifyAll('how',123,'asoenhuas'), False)
        self.assertTrue(foo('how',123))
        self.assertTrue(foo(abc='how',xyz=123,big=None))
        
if __name__ == '__main__':
    
    unittest.main()