Add return annotation support to u.quantity_input#5606
Conversation
| with add_enabled_equivalencies(self.equivalencies): | ||
| return wrapped_function(*func_args, **func_kwargs) | ||
| return_ = wrapped_function(*func_args, **func_kwargs) | ||
| if wrapped_signature.return_annotation is not funcsigs._empty: |
There was a problem hiding this comment.
Is it safe to rely on _empty since it's private?
There was a problem hiding this comment.
https://docs.python.org/3/library/inspect.html suggests that there is a non-private Signature.empty property. Similarly, for funcsigs: http://funcsigs.readthedocs.io/en/0.4/
|
|
||
| import astropy.units as u | ||
| @u.quantity_input | ||
| def myfunction(myangle: u.arcsec) -> u.deg: |
There was a problem hiding this comment.
This example would always feel, wouldn't it? It would need return units of u.deg**2.
|
Looks good modulo a few comments. Nice to see how relatively easy this is! |
| with add_enabled_equivalencies(self.equivalencies): | ||
| return wrapped_function(*func_args, **func_kwargs) | ||
| return_ = wrapped_function(*func_args, **func_kwargs) | ||
| if isinstance(wrapped_signature.return_annotation, (Unit, Quantity)): |
There was a problem hiding this comment.
@astrofrog @mhvk I changed this to verify Unit or Quantity so that quantity_input doesn't just eat all the possible uses for return annotations. Opinions?
There was a problem hiding this comment.
Probably a good idea, though is there a use for Quantity here? Why not just Unit?
There was a problem hiding this comment.
I don't know why you would but:
❯ (10*u.m).to(100*u.km)
<Quantity 0.0001 100 km>is valid...
There was a problem hiding this comment.
I have a slight worry here, because function units such as u.mag(u.Jy) or u.STmag are not Unit instances (this does need to be fixed eventually; #3650). My sense is that if one uses @quantity_input, then one should not really expect to be able to use the output annotation for anything else, so I think I would just stick with what you had. But if you do want to be more general, I would suggest:
try:
return_unit = Unit(wrapped_signature.return_annotation)
except Exception: # not a unit requested, so ignore
pass
else:
return return_.to(return_unit)
Obviously, the above, like your original implementation, allows one to provide units as strings, which I think is OK.
There was a problem hiding this comment.
I have reverted it to assuming the return annotation is always a Unit-Like thing. #3650 looks quite interesting...
|
Also, I just learned that as of Python 3.2 the function annotations are in a |
cb33e85 to
e634892
Compare
|
@Cadair - this is really neat, and I think it deserves a mention in the narrative docs as well... Perhaps either the quantities page or even the units getting started page? Also, when this gets released, I think it makes sense to integrate it into the Quantities Tutorial, which has a section on using quantities in functions. |
|
I like. Might make me finally upgrade from 2.7 to 3.x! |
|
@eteq I added an example to the existing doc section on functions, it seemed like the best place. |
|
@astrofrog - This looks good, and since you already approved, I'll go ahead and merge. Thanks, @Cadair! |
|
Things like this are definitely an incentive for me to switch my production environment to python3 ASAP.Thanks! |
|
@hamogu - Another good reason is that py3 is significantly faster... |
|
My main problem is that mayavi is not fully ported yet and I use that in some projects, but I'm going off-topic now... |
|
@hamogu there is now a linux mayavi / vtk build on conda-forge for python 3 and linux |

This adds a lovely Python 3 only feature to
u.quantity_inputwhich enables the use of the return annotation to specify a unit to convert the output to. This makes it much easier to unit convert the return value of a function, rather than explicitly doing it in the return statement. I have found this especially useful for when a value is often printed to screen.I have assumed here that the error handling in
Quantity.tois sufficient for type checking on the return value itself, so this does not need to be done in the decorator.Also, this could be hacked to be 2.7 compatible by a special keyword argument to
quantity_inputbut frankly that's ugly and I can't be bothered...