Writing unit aware functions that accept pint.Quantity and string arguments

[jrmagers]

I like to use Pint to write functions for unit-aware calculations:

import pint

ureg = pint.UnitRegistry()

def rectangle_area(length: pint.Quantity, width: pint.Quantity) -> pint.Quantity:
    """Compute area of a rectangle."""

    area = length * width
    return area.to_base_units()

rectangle_area(length=5 * ureg.cm, width=3 * ureg.inch)

While the above works great, I would also like to use the same function but also have it accept string arguments as well. This is mostly a matter of convenience.

rectangle_area(length="5 cm", width="3 inch")

Is there a recommended way to achieve this?

[jrmagers]

Here's a solution:

import pint

ureg = pint.UnitRegistry()
ureg.formatter.default_format = "~"  # display abbreviated units

def typecast_string_args(*args):
    """Cast string arguments to pint.Quantity."""

    casted_args = []

    for arg in args:
        if isinstance(arg,str):
            casted_args.append(ureg(arg))
        elif isinstance(arg,pint.Quantity):
            casted_args.append(arg)
        else:
            raise TypeError("Arguments must be of type str or pint.Quantity.")
    
    return casted_args

def rectangle_area(length: pint.Quantity | str, width: pint.Quantity | str) -> pint.Quantity:
    """Compute area of a rectange."""

    length, width = typecast_string_args(length, width)

    area = length * width
    return area.to_base_units()

Works when all arguments are strings, all arguments are pint.Quantity or mixed:

rectangle_area(length=50 * ureg.cm, width=100 * ureg.inch)
rectangle_area(length="50 cm", width="100 inch")
rectangle_area(length=50 * ureg.cm, width="100 inch")