manifold/util/type.py

   1 # http://wiki.python.org/moin/PythonDecoratorLibrary#Type_Enforcement_.28accepts.2Freturns.29
   2 '''
   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.
   9
  10 Example usage:
  11     >>> NONE, MEDIUM, STRONG = 0, 1, 2
  12     >>>
  13     >>> @accepts(int, int, int)
  14     ... @returns(float)
  15     ... def average(x, y, z):
  16     ...     return (x + y + z) / 2
  17     ...
  18     >>> average(5.5, 10, 15.0)
  19     TypeWarning:  'average' method accepts (int, int, int), but was given
  20     (float, int, float)
  21     15.25
  22     >>> average(5, 10, 15)
  23     TypeWarning:  'average' method returns (float), but result is (int)
  24     15
  25
  26 Needed to cast params as floats in function def (or simply divide by 2.0).
  27
  28     >>> TYPE_CHECK = STRONG
  29     >>> @accepts(int, debug=TYPE_CHECK)
  30     ... @returns(int, debug=TYPE_CHECK)
  31     ... def fib(n):
  32     ...     if n in (0, 1): return n
  33     ...     return fib(n-1) + fib(n-2)
  34     ...
  35     >>> fib(5.3)
  36     Traceback (most recent call last):
  37       ...
  38     TypeError: 'fib' method accepts (int), but was given (float)
  39
  40 '''
  41 import sys
  42 from itertools import izip
  43
  44 def accepts(*types, **kw):
  45     '''Function decorator. Checks decorated function's arguments are
  46     of the expected types.
  47
  48     Parameters:
  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).
  53              debug = ( 0 | 1 | 2 )
  54
  55     '''
  56     if not kw:
  57         # default level: MEDIUM
  58         debug = 2
  59     else:
  60         debug = kw['debug']
  61     try:
  62         def decorator(f):
  63             # XXX Missing full support of kwargs
  64             def newf(*args, **kwargs):
  65                 if debug is 0:
  66                     return f(*args, **kwargs)
  67                 assert len(args) == len(types)
  68                 argtypes = tuple(map(type, args))
  69                 if not compare_types(types, argtypes):
  70                 # if argtypes != types:
  71                     msg = info(f.__name__, types, argtypes, 0)
  72                     if debug is 1:
  73                         print >> sys.stderr, 'TypeWarning: ', msg
  74                     elif debug is 2:
  75                         raise TypeError, msg
  76                 return f(*args, **kwargs)
  77             newf.__name__ = f.__name__
  78             return newf
  79         return decorator
  80     except KeyError, key:
  81         raise KeyError, key + "is not a valid keyword argument"
  82     except TypeError, msg:
  83         raise TypeError, msg
  84
  85 def compare_types(expected, actual):
  86     if isinstance(expected, tuple):
  87         if isinstance(actual, tuple):
  88             for x, y in izip(expected, actual):
  89                 if not compare_types(x ,y):
  90                     return False
  91             return True
  92         else:
  93             return actual == type(None) or actual in expected
  94     else:
  95         return actual == type(None) or actual == expected or isinstance(actual, expected) # issubclass(actual, expected)
  96
  97 def returns(ret_type, **kw):
  98     '''Function decorator. Checks decorated function's return value
  99     is of the expected type.
 100
 101     Parameters:
 102     ret_type -- The expected type of the decorated function's return value.
 103                 Must specify type for each parameter.
 104     kw       -- Optional specification of 'debug' level (this is the only valid
 105                 keyword argument, no other should be given).
 106                 debug=(0 | 1 | 2)
 107     '''
 108     try:
 109         if not kw:
 110             # default level: MEDIUM
 111             debug = 1
 112         else:
 113             debug = kw['debug']
 114         def decorator(f):
 115             def newf(*args):
 116                 result = f(*args)
 117                 if debug is 0:
 118                     return result
 119                 res_type = type(result)
 120                 if not compare_types(ret_type, res_type):
 121                 # if res_type != ret_type: # JORDAN: fix to allow for # StringTypes = (str, unicode)
 122                 # XXX note that this check should be recursive
 123                     msg = info(f.__name__, (ret_type,), (res_type,), 1)
 124                     if debug is 1:
 125                         print >> sys.stderr, 'TypeWarning: ', msg
 126                     elif debug is 2:
 127                         raise TypeError, msg
 128                 return result
 129             newf.__name__ = f.__name__
 130             return newf
 131         return decorator
 132     except KeyError, key:
 133         raise KeyError, key + "is not a valid keyword argument"
 134     except TypeError, msg:
 135         raise TypeError, msg
 136
 137 def info(fname, expected, actual, flag):
 138     '''Convenience function returns nicely formatted error/warning msg.'''
 139     format = lambda types: ', '.join([str(t).split("'")[1] for t in types])
 140     msg = "'{}' method ".format( fname )\
 141           + ("accepts", "returns")[flag] + " ({}), but ".format(expected)\
 142           + ("was given", "result is")[flag] + " ({})".format(actual)
 143     return msg
 144