Skip to content

Instantly share code, notes, and snippets.

@haydenbbickerton
Created April 19, 2017 01:42
Show Gist options
  • Save haydenbbickerton/5a968a5e972ff7dc8718db22b9714b3e to your computer and use it in GitHub Desktop.
Save haydenbbickerton/5a968a5e972ff7dc8718db22b9714b3e to your computer and use it in GitHub Desktop.
Annotations decorator for python 2.7.
def annotate(*args, **kwargs):
'''
Decorator to add PEP 3107 style function annotations, and enforce the types
if desired. Annotations are added as the "__annotations__" attribute.
Can read more here - http://python-future.org/func_annotations.html
Parameters
----------
enforce : bool
If True, the parameters will be checked against the expected types
before running (the default is False)
returns : mixed
The type of object returned by function call (the default is None)
**types : mixed
The keyword arguments and their expected types. Can be a single type
or a list of acceptable types. Can also be a string, in which case
comparison will be done against the type().__name__ attribute.
Examples
----------
> @annotate(int, int, returns=int)
> def my_add(a, b):
> return a + b
> print my_add.__annotations__ # {'a': int, 'b': int, 'returns': int}
>
...
> @annotate(a=int, b=int, returns=int, enforce=True)
> def my_add(a, b):
> return a + b
>
> my_add(123, '456') # TypeError: Expected type of "int" for parameter "b",
> recieved "str" instead.
>
...
> @annotate(a=int, b=(int, float), c='MyCustomClass', returns=int, enforce=True)
> def my_add(a, b, c, d='fus ro dah'):
> # We didn't specify a type for d, so it's ignored
> return int(a + b)
> print my_add.__annotations__ # {'a': int, 'b': (int, float),
'c': 'MyCustomClass', 'returns': int}
'''
enforce = kwargs.pop('enforce', False)
returns = kwargs.pop('returns', None)
def outer(func):
types = dict(zip(func.__code__.co_varnames, args))
types.update(kwargs)
def inner(func, *args, **kwargs):
if enforce is True:
params = dict(zip(func.__code__.co_varnames, args))
params.update(kwargs)
# Only enforce for params specified
check_types = {k: v for k, v in params.items() if k in types}
_name = lambda x: getattr(x, '__name__', x)
_check = lambda x, y: _name(type(x)) == y if isinstance(y, basestring) else isinstance(x, y)
# Ensure passed types are the same as expected types
for param, value in check_types.items():
_types = types[param]
_types = (_types,) if isinstance(_types, type) else _types
valid = any(_check(value, _type) for _type in _types)
if not valid:
raise TypeError(
'Expected type of "{}" for parameter "{}", recieved type of "{}" instead.'
.format(', '.join([_name(x) for x in _types]), param, type(value).__name__)
)
return func(*args, **kwargs)
func.__annotations__ = types
func.__annotations__['returns'] = returns
return decorator.decorator(inner, func)
return outer
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment