Decorator to convert a function's result into the desired units

I tried to collect some of my formulas inside classes so that I didn't need to browser the names of the functions anymore but only determine which Quantity I want to calculate and then check which methods are avaiable for determining it. This is based on the astropy.units module. Just to make a small example:

import astropy.units as u

class Length(object):
unit = u.m

@classmethod
#@quantitydecorator # I'll explain this later
def from_velocity_time(cls, velocity, time):
"""
Calculate length from a given velocity and time.

Parameters
----------
velocity: astropy.units.Quantity
The velocity of the object.
time: astropy.units.Quantity
The time the object moved at this velocity

Returns
-------
length: astropy.units.Quantity
The determined length based on velocity * time
"""
return velocity * time


This works great, I tried:

Length.from_velocity_time(100*u.km/u.h, 100*u.h)
# gives 10000 km


but has some downsides when the units are equivalent but different:

Length.from_velocity_time(100*u.km/u.s, 100*u.h)
# 10000 h*km/s


so I decided to create a general decorator that converts the results unit to the unit defined in the class or have an additional parameter result_unit for every method call:

def quantitydecorator(f):
"""
Allows for an extra argument "result_unit" for all methods
and converts the result to the unit specified there or if not
given to the unit specified in the class.
"""
@wraps(f)
def wrapper(*args, **kwds):
# Check if result_unit is given otherwise use the class default
# (args[0] is the class)
if 'result_unit' in kwds:
result_unit = u.Unit(kwds.pop('result_unit'))
else:
result_unit = args[0].unit

# Calculate the result
result = f(*args, **kwds)

# Convert to result unit if the result has a unit
if not hasattr(result, 'unit'):
# No unit attribute so we have a plain numpy array
# or number, only let it pass when target unit is
# dimensionless
if args[0].unit != u.dimensionless_unscaled:
raise ValueError('Got dimensionless quantity but needed'
' quantity with unit: {0}'.format(args[0].unit))
# Result has a different unit than wanted, convert it to wanted unit
elif result.unit != result_unit:
result = result.to(result_unit)
return result
return wrapper


This works fine (this time the decorator must be uncommented in the class):

Length.from_velocity_time(100*u.km/u.s, 100*u.h)
# 3.6e10m

Length.from_velocity_time(100*u.km/u.s, 100*u.h, result_unit=u.AU)
# 0.24064514AU


but I can't help suspecting that this is rather awkward. I'm using approximately 300 functions but there are only around 30-40 parameters (like Length in the example but more sophisticated like "free-fall-time", "half-light-radius". "redshift", etc. everything that I at some point needed for any calculation and there are typically 2-10 ways to calculate such a Quantity) so it would be a great gain to collect them into classes and don't have duplicate logic with the result_unit. Would you consider this "good coding" style or should I dismiss this attempt and choose a different approach (if so, any idea which)?

Length.from_velocity_time(100*u.km/u.s, 100*u.h, result_unit=u.AU)


… is less conventional and not any better than

Length.from_velocity_time(100*u.km/u.s, 100*u.h).to(u.AU)


In addition, monkey-patching the from_velocity_time() function to support a special result_unit named parameter feels slimy.

The only drawback with .to(…) is that pure numbers wouldn't support such a method. But that would mean that the calculation is complete nonsense anyway.

• Thank for your input. After some consideration I think you are completly right. I was overthinking this and just a bit too annoyed especially if constants are involved that the result often contained km / m stuff. I'll stick with the .to method and discard that decorator. Many thanks. Feb 24, 2016 at 0:01

I agree with everything said in 200_success's answer. It also seems like maybe instead of passing in a monkey-patched keyword argument for the desired unit, to instead do some automatic unit simplification with some simplification keyword argument that requires no decoration, and then allow the user to transform to whatever units they want it in if they need to.

Overall, it seems like you're doing extra computational work to meet a need the user doesn't necessarily have, but is perfectly capable of doing on their own.

I'd also like to address a minor code smell I see in what you've written.

You seem to always ask permission before you do things (i.e. hasattr, x in l, etc). This isn't necessarily bad, but generally in Python we prefer to beg forgiveness if we need to. In particular, and for Python 2 especially, hasattr is known to cause problems. I'd rewrite your wrapper function to look like this.

result_unit = u.Unit(kwds.pop('result_unit', args[0].unit))

result = f(*args, **kwds)

try:
if result.unit != result_unit:
result = result.to(result_unit)
except AttributeError:
# No unit attribute so we have a plain numpy array
# or number, only let it pass when target unit is
# dimensionless
if result_unit != u.dimensionless_unscaled:
raise ValueError('Got dimensionless quantity but needed'
' quantity with unit: {0}'.format(args[0].unit))
return result


You also (seemingly) have a small bug - you said "only let it pass when target unit is dimensionless", but then compared args[0].unit instead of result_unit, which seems like it won't always be correct.