from __future__ import print_function
import simtk.unit as units
[docs]class UnitsException(Exception):
"""Exception denoting that an argument has the incorrect units."""
def __init__(self, value):
self.value = value
def __str__(self):
return repr(self.value)
[docs]class ValueException(Exception):
"""Exception denoting that an argument has the incorrect value."""
def __init__(self, value):
self.value = value
def __str__(self):
return repr(self.value)
[docs]def accepts(*types):
"""
Decorator for class methods that should accept only specified types.
EXAMPLE
@accepts(float, int)
def function(a, b):
return b*a
"""
def check_accepts(f):
nargs = (f.func_code.co_argcount - 1) # exclude self
assert len(types) == nargs, ("Incorrect number of args supplied in "
"@accepts decorator for class method %s" % (f.func_name))
def new_f(*args, **kwds):
for (a, t) in zip(args[1:], types):
if a is not None:
assert isinstance(a, t), "arg %r does not match %s" % (a, t)
return f(*args, **kwds)
new_f.func_name = f.func_name # copy function name
new_f.func_doc = f.func_doc # copy docstring
return new_f
return check_accepts
[docs]def accepts_compatible_units(*units, **unitdict):
"""
Decorator for class methods that should accept only arguments compatible with specified units.
Each argument of the function will be matched with an argument of @acceptunits.
Those arguments of the function that correspond @acceptunits which are not None
will be checked to ensure they are compatible with the specified units.
EXAMPLE
@acceptsunits(units.meter, None, units.kilocalories_per_mole)
def function(a, b, c): pass
function(1.0 * units.angstrom, 3, 1.0 * units.kilojoules_per_mole)
"""
def check_units(f):
nargs = (f.func_code.co_argcount - 1) # exclude self
#assert len(units) == nargs, "incorrect number of units supplied in @accepts_compatible_units decorator for class method %s" % (f.func_name)
def new_f(*args, **kwds):
for (a, u) in zip(args[1:], units):
if u is not None:
assert (a.unit).is_compatible(u), "arg %r does not have units compatible with %s" % (a,u)
for (k, u) in unitdict.iteritems(): # should be ignored if no unitdict
if u is not None:
assert (kwds[k].unit).is_compatible(u), "kwd arg %s does not have units compatible with %s" % (kwds[k],u)
return f(*args, **kwds)
new_f.func_name = f.func_name # copy function name
new_f.func_doc = f.func_doc # copy docstring
return new_f
return check_units
[docs]def returns(rtype):
"""
Decorator for functions that should only return specific types.
EXAMPLE
@returns(int)
def function(): return 7
"""
def check_returns(f):
def new_f(*args, **kwds):
result = f(*args, **kwds)
assert isinstance(result, rtype), "return value %r does not match %s" % (result,rtype)
return result
new_f.func_name = f.func_name # copy function name
new_f.func_doc = f.func_doc # copy docstring
return new_f
return check_returns