1 # http://wiki.python.org/moin/PythonDecoratorLibrary#Type_Enforcement_.28accepts.2Freturns.29
3 One of three degrees of enforcement may be specified by passing
4 the 'debug' keyword argument to the decorator:
5 0 -- NONE: No type-checking. Decorators disabled.
6 1 -- MEDIUM: Print warning message to stderr. (Default)
7 2 -- STRONG: Raise TypeError with message.
8 If 'debug' is not passed to the decorator, the default level is used.
11 >>> NONE, MEDIUM, STRONG = 0, 1, 2
13 >>> @accepts(int, int, int)
15 ... def average(x, y, z):
16 ... return (x + y + z) / 2
18 >>> average(5.5, 10, 15.0)
19 TypeWarning: 'average' method accepts (int, int, int), but was given
22 >>> average(5, 10, 15)
23 TypeWarning: 'average' method returns (float), but result is (int)
26 Needed to cast params as floats in function def (or simply divide by 2.0).
28 >>> TYPE_CHECK = STRONG
29 >>> @accepts(int, debug=TYPE_CHECK)
30 ... @returns(int, debug=TYPE_CHECK)
32 ... if n in (0, 1): return n
33 ... return fib(n-1) + fib(n-2)
36 Traceback (most recent call last):
38 TypeError: 'fib' method accepts (int), but was given (float)
42 from itertools import izip
44 def accepts(*types, **kw):
45 '''Function decorator. Checks decorated function's arguments are
46 of the expected types.
49 types -- The expected types of the inputs to the decorated function.
50 Must specify type for each parameter.
51 kw -- Optional specification of 'debug' level (this is the only valid
52 keyword argument, no other should be given).
57 # default level: MEDIUM
66 assert len(args) == len(types)
67 argtypes = tuple(map(type, args))
68 if not compare_types(types, argtypes):
69 # if argtypes != types:
70 msg = info(f.__name__, types, argtypes, 0)
72 print >> sys.stderr, 'TypeWarning: ', msg
76 newf.__name__ = f.__name__
80 raise KeyError, key + "is not a valid keyword argument"
81 except TypeError, msg:
84 def compare_types(expected, actual):
85 if isinstance(expected, tuple):
86 if isinstance(actual, tuple):
87 for x, y in izip(expected, actual):
88 if not compare_types(x ,y):
92 return actual == type(None) or actual in expected
94 return actual == type(None) or actual == expected or issubclass(actual, expected)
96 def returns(ret_type, **kw):
97 '''Function decorator. Checks decorated function's return value
98 is of the expected type.
101 ret_type -- The expected type of the decorated function's return value.
102 Must specify type for each parameter.
103 kw -- Optional specification of 'debug' level (this is the only valid
104 keyword argument, no other should be given).
109 # default level: MEDIUM
118 res_type = type(result)
119 if not compare_types(ret_type, res_type):
120 # if res_type != ret_type: # JORDAN: fix to allow for # StringTypes = (str, unicode)
121 # XXX note that this check should be recursive
122 msg = info(f.__name__, (ret_type,), (res_type,), 1)
124 print >> sys.stderr, 'TypeWarning: ', msg
128 newf.__name__ = f.__name__
131 except KeyError, key:
132 raise KeyError, key + "is not a valid keyword argument"
133 except TypeError, msg:
136 def info(fname, expected, actual, flag):
137 '''Convenience function returns nicely formatted error/warning msg.'''
138 format = lambda types: ', '.join([str(t).split("'")[1] for t in types])
139 msg = "'{}' method ".format( fname )\
140 + ("accepts", "returns")[flag] + " ({}), but ".format(expected)\
141 + ("was given", "result is")[flag] + " ({})".format(actual)