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

Question

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)?

Answer 1

Your proposal…

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.

Answer 2

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.