Skip to content

Instantly share code, notes, and snippets.

@cathalgarvey
Created May 1, 2014 14:30
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save cathalgarvey/aa376981d9861935a048 to your computer and use it in GitHub Desktop.
Save cathalgarvey/aa376981d9861935a048 to your computer and use it in GitHub Desktop.
A decorator for when you want strictness in Python3.
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