Type hinting issue -> stop using `pint.wraps`
The type hinting with Pint's wraps functionality can't work because the decorator is only created at runtime so you obviously can't do static type hinting with it. To solve this, verify_units just returns a function with the same signature as is passed in. That's fine, but I had a thought. Given we only really want to use wraps to convert Quantity to their underlying type (float, int or array) and then back, I think it would be clearer to just do that inside the function rather than using verify_units. That would also allow us to give the user control so they can pass in the unit registry at runtime if they wish (falling back to get_application_registry otherwise) and remove a layer of misdirection, please remove reliance on one of pint's more confusing (and impossible to type hint properly) APIs.
@lewisjarednz I'm interested in whether you have any strong thoughts about why the pattern below wouldn't work instead of anything that currently relies on wraps / verify_units?
I have confirmed by hand that this refactoring would fix my downstream mypy issues I'm having in the copier template (although now having stared at this for a while I think the current setup would probably work too, but now that I think I've seen this simplification and how to avoid the verify_units workaround I would like to make the change for easier long-term maintenance).
def calc_flux(self, param: Quantity, ur: pint.registry.UnitRegistry | None = None) -> Quantity:
# strip untis before passing to wrapper Fortran modules
param = param.to(_UNITS["param"]).m
# call Fortran
out_raw = wrapped.calc_flux(self.id, out_raw)
# put units back on before returning
# use unit registry of users' choice, otherwise default application registry
if ur is None:
ur = pint.get_application_registry()
return ur.Quantity(out_raw, _UNITS["flux"])