Created
May 1, 2014 14:30
-
-
Save cathalgarvey/aa376981d9861935a048 to your computer and use it in GitHub Desktop.
A decorator for when you want strictness in Python3.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
def strict(func): | |
'''A decorator for methods or functions that requires annotations for all | |
arguments and the return value, throws typeerrors on deviations. | |
Remember that for more than one return value, the return type is "tuple". | |
Container-type arguments or return values are only inspected at top-level. | |
Note that as written, this does not handle catchall argument types "*args", or "**kwargs". | |
''' | |
import inspect, collections | |
NoneType = type(None) | |
def die_on_untyped_annotation(par, ann_type="argument annotation"): | |
'Saves some lines/clarity by functionalising checks and Exception raising.' | |
if par == inspect._empty: | |
raise ValueError(("Error decorating strictly typed function '{}': " | |
"{} not provided.").format(func, ann_type)) | |
argn = getattr(par, 'name', 'return') | |
try: | |
ann = par.annotation | |
except AttributeError: | |
ann = par.return_annotation | |
if ann == None: | |
ann = NoneType | |
if isinstance(ann, tuple): | |
for a in ann: | |
if a == None: | |
a = NoneType | |
if not isinstance(a, type): | |
raise ValueError(("Error decorating strictly typed function '{}': " | |
"{} for {} '{}' is not a class.").format(func, ann_type, argn, ann)) | |
elif not isinstance(ann, type): | |
raise ValueError(("Error decorating strictly typed function '{}': " | |
"{} for {} '{}' is not a class.").format(func, ann_type, argn, ann)) | |
# == Assess annotations and copy relevant signature bits == | |
# At this point, if any args/returns are None, change to NoneType. | |
funcsig = inspect.signature(func) | |
funcreturn = funcsig.return_annotation | |
if funcreturn == None: | |
funcreturn = NoneType | |
funcargs = funcsig.parameters.copy() | |
for k in funcargs: | |
if funcargs[k].annotation == None: | |
funcargs[k].annotation = NoneType | |
# Before decorating, make sure annotations and return value are classes? | |
for arg_annotation in funcargs.values(): | |
die_on_untyped_annotation(arg_annotation) | |
die_on_untyped_annotation(funcsig, "return annotation") | |
# Decorated function. | |
def new_func(*args, **kwargs): | |
# Check positional args | |
for num, arg, funcarg in zip(range(1,len(args)+1), args, funcargs.values()): | |
f_ann = funcarg.annotation | |
if not isinstance(arg, f_ann): | |
raise TypeError("Argument of wrong type (expected {}) passed to this function as argument {} with type {}: {}".format(f_ann, num, type(arg), arg)) | |
# Check keyword args | |
for arg, value in kwargs.items(): | |
if arg in funcargs: | |
f_ann = funcargs[arg].annotation | |
if not isinstance(value, f_ann): | |
raise TypeError("Argument of wrong type (expected {}) passed to this function as keyword argument {} with type {}: {}".format(f_ann, arg, type(value), value)) | |
# Run function | |
retval = func(*args, **kwargs) | |
# Check return value | |
if not isinstance(retval, funcreturn): | |
raise TypeError("Error: Return value of this function is of wrong type; expected {}, got type {}: {}".format(funcreturn, type(retval), retval)) | |
return retval | |
# Keep documentation of old function, plus extra line indicating strict typing. | |
additional_docs = ''' (This function is strictly typed: {}({})->{})'''.format( | |
func.__name__, ', '.join(('{}:{}'.format(k,v.annotation) for k,v in funcargs.items())), funcreturn) | |
new_func.__doc__ = (func.__doc__ or '') + additional_docs | |
return new_func |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment