Allow return annotations other than Quantity

Description

The decorator quantity_input() tries to convert the annotated return value to a Quantity here. When a return annotation other than Quantity is used, this conversion fails of course resulting in the error below.

Expected behavior

The decorator quantity_input() should ignore annotated return values which are not of type Quantity. This has also been disicussed and fixed here but the fix has been discarded.

Actual behavior

The following error message is shown:

Traceback (most recent call last):
  File "<input>", line 6, in <module>
  File "<<redacted>>/venv/lib/python3.7/site-packages/astropy/units/decorators.py", line 236, in wrapper
    return return_.to(wrapped_signature.return_annotation)
AttributeError: 'str' object has no attribute 'to'

Steps to Reproduce

[Run the code below in a python console]

import astropy.units as u

@u.quantity_input(temp=[u.Kelvin, u.Celsius])
def foo(temp: u.Quantity) -> str:
    return "bar"
print(foo(300 * u.K)())

System Details

Darwin-19.4.0-x86_64-i386-64bit
Python 3.7.3 (v3.7.3:ef4ec6ed12, Mar 25 2019, 16:52:21) 
[Clang 6.0 (clang-600.0.57)]
Numpy 1.18.2
astropy 4.0.1.post1
Scipy 1.4.1

👍 To supporting this.

The usecase under this, of coming up with a way of supporting type annotations and unit validation at the same time is interesting, and not one the current implementation of quantity_input was built for. This is the first example of a thing that makes me want to keep they keyword argument syntax to quantity_input around for reasons other than backwards compatibility. (Up to this point I would have liked to deprecate it in favour of function annotations).

I do wonder if there is something we could provide which is a unit aware version of a Quantity type annotation, i.e. 'Quantity[u.Kelvin]' or similar. My current experience with the type hinting support in Python in limited so this probably needs someone who knows more about it in the kind of contexts it would be useful.

I've been running into this recently and it would be very nice for quantity_input to support both type annotations and unit validation. Having quantity_input ignore anything which is not a Quantity or UnitBase input would be part of the solution.

I really like the suggestion for a unit aware version of a Quantity type annotation. One possibility is the PEP to add a type annotation wrapper for metadata. As example,

import astropy.units as u
from typing import Annotated  # not yet in library, see PEP

@u.quantity_input()
def foo(temp: Annotated[u.Quantity, u.Kelvin]):
    return "bar"

Another option, to have Quantity provide the type, might be to use typing.TypeVar(..., bound=Quantity). However, I have yet to find a way to actually tell TypeVar that the variable is a Quantity type but with some specific unit.

I agree that ideally we move towards supporting some form of Quantity[<unit>]. There has been a general problem with type annotations for arrays, but I noticed lots of numpy is being annotated now. See https://github.com/numpy/numpy/issues?q=is%3Aopen+sort%3Aupdated-desc+label%3A%22static+typing%22

Though I think that long-term plan shouldn't stop us from a short-term solution of allowing something else than units for the return type!

Are there plans to static type Astropy? or at least add type hints like CoordinateType or UnitType, etc?
I also have a fair number of types I created for personal use, like UnitableType which allows anything that Unit can accept as an argument, like UnitBase instances, Quantity instances, or strings. If there were an AEP or PR I'd be happy to chip in where I can.

typing_extensions is the official backport for Annotated and works on any py3.5+. I suggest that this be the way in which quantity_input stores and determines all unit information. The __annotations__ for the wrapper function which quantity_input produces would store all this information so that the original function would be unchanged. Since decorators can be applied as functions, not just with the "pie" syntax, this would allow multiple different versions of the same function to be created with different unit specifications. Not only that, but storing all the information in __annotations__ means that the desired unit outputs are pre-computed, increasing function call speed, and are also editable by knowledgeable users in a python-standard format.

The other advantage of this method is that statically typed libraries would be able to also support quantity_input by using Annotated[<type>, <quantity_input_argument>]. This could even be a temporary implementation to the Quantity[<unit>] idea since Quantity.__getitem__ could return Annotated[Quantity, <unit>].

Some thoughts, on an example function

@quantity_input(y="speed", v=[u.km, u.parsec])
def foo(x: u.s, y: u.Quantity, v=2 * u.m) -> "distance":
     return x * v + y

# annotations
foo.__annotations__ == {"x": Annotated[Quantity, u.s],
                                       "y": Annotated[Quantity, "speed"],
                                       "v": Annotated[Quantity, [u.km, u.parsec]],
                                       "return_": Annotated[Quantity, "distance"]}
#and the original function
foo.__wrapped__.__annotations__ == {"x": u.s, "y": Quantity, "return_": "distance"}

Order of precedence:

Kwargs into quantity_input:
quantity_input(x=u.hour)(foo) would convert x in foo to hours not seconds.
This allows quantity_input to wrap functions from external libraries without modifying the function.
function annotations
So that calling quantity_input as a function overrides the default unit specifications.
units from argument default values
To respect type hints if Quantity[<unit>] is implemented

There needs to be an override such that an argument can be skipped. I think kwarg with value None is the simplest flag. Eg,

@quantity_input(return_=None)
def foo(x) -> "distance":
     return x

# annotations
foo.__annotations__ == {"x": Annotated[Any, None],
                                       "return_": Annotated[Any, None]}

Since Annotated supports multiple annotations, we could even preserve the original annotations in the __wrapped__ function by just shifting them to the second slot in Annotated and reserve the first slot for quantity_input arguments. This would make annotation comments (like x: "the time from event 1 to 2") not be hidden.

@mhvk

@mhvk @Cadair, I solved the Quantity[<unit>] -> Annotated[Quantity, unit, other_annotations]

def _parse_allowed_type_hint(target, detailed_exception=True):
    try:  # unit passed in as a string or unit or Quantity
        target_unit = Unit(target)

    except ValueError:  # physical type or mistake
        from astropy.units.physical import _unit_physical_mapping  # TODO move import
        if target in _unit_physical_mapping:
            target_unit = target
        else:   # Function argument target is invalid
            if detailed_exception:
                raise ValueError(
                    (
                        f"Invalid unit or physical type '{target}'.\n"
                        f"Did you mean: {'TODO!'}"
                        # TODO did_you_mean(target, candidates)
                    )
                )
            else:
                raise ValueError(f"Invalid unit or physical type '{target}'.")

    return target_unit

class Quantity:
    def __class_getitem__(cls, unit_and_annotations):
        if isinstance(unit_and_annotations, tuple):  # unit and annotations
            unit = unit_and_annotations[0]
            if len(unit_and_annotations) > 1:
                annotations = unit_and_annotations[1:]
            else:  # just unit
                annotations = ()
        else:  # just unit
            unit = unit_and_annotations
            annotations = ()

        # check unit / physical type is allowed
        unit = _parse_allowed_type_hint(unit, detailed_exception=True)

        return Annotated.__class_getitem__((cls, UnitSpec(unit, dtype=cls), *annotations))

Note: I've edited this to include UnitSpec, see #10156 (comment)

@nstarman

@nstarman - I've not looked into type hints in any detail myself (and cannot find much about Annotate - do you have a link?), but your example looks very nice indeed! @Cadair, what do you think?
In case it makes things easier: for 4.2 we'll likely only support python >=3.7, since 3.6 is reaching end of support in December.

The relevant PEP fully describing Annotated is https://www.python.org/dev/peps/pep-0593/#id13. In brief, it is a wrapper for a static type that allows for attachment of metadata, ie Annotated[static_type, metadata, metadata,...]. Since quantity_input always has "static_type" as a Quantity (or subclass) this is easy. We put the "quantity_helper" unit info to the first metadata field, moving any existing metadata.

Unfortunately Annotated will be core Python in 3.9+, so we'll be on the official backport for a while

Thanks, that indeed looks like just what we need! I had worried that our clever use of the annotations was going to cause problems with typing (but when we implemented it, it wasn't as clear that was what would come to dominate; great to see it addressed in PEP 593).

I tried mocking up an implementation for creating the __annotations__ dictionary for the wrapper. It mostly works, but the number of options / if statements / etc to support static typing is pretty high. For instance, detecting a user-created Annotated[Quantity, "comment"] is not aAnnotated[Quantity, <unit or physical_type>, "comment"] is both hard and needlessly expensive to do at each function evaluation.

@quantity_input
def func(x: Annotated[Quantity, "a general comment"]):
    pass

The ideal solution is to subclass Annotated to make a QuantityAnnotated type, but this is forbidden by Annotated.__init_subclass__. The next-best thing is to create some lightweight wrapper class — maybe called UnitSpec s.t. the function evaluations need only check if the first piece of metadata is that wrapper. In pseudocode (pretending quantity_input is a function-type, not class-type, decorator factory)

def quantity_input(func=None, equivalencies=[], **quantity_annotations):
    if func is None:  return functools.partial(quantity_input, equivalencies= equivalencies, **quantity_annotations)

    # Make __annotations__
    sig = inspect.signature(func)
    annotations = dict()
    for param in sig.parameters.values():
        ...  # do stuff here to make the annotations given: kwargs, type annotations, and default values
       # eg    dict(x=Annotated[Quantity, UnitSpec("distance")], y=str)

    @wraps(func)
    def wrapper(*args, **kwargs):
        ba = sig.bind_partial(*args, **kwargs)
        ba.apply_defaults()

        with add_enabled_equivalencies(equivalencies):
       
            for name, annot in wrapper.__annotations__.items():  # this updates with external changes to __annotations__ :)!
                if annot is Annotated:
                    unitspec = annot.__metadata__[0]
                    if isinstance(unitspec, UnitSpec):
                        ba.arguments[name] = unitspec(ba.arguments[name], dtype=annot.__origin__)  # cast to right type
                    else: pass  # not actually a unitspec
               else: pass  # just a normal annotation

            return func(*ba.args, **ba.kwargs)  # call with equivalencies
    # /def

    wrapper.__annotations__ = annotations  # set annotations
    wrapper._equivalencies = equivalencies  # store equivalencies
    return wrapper

Of course, I've neglected to do the return_annotation conversion, but that's easily implemented.

I've imagined here that UnitSpec has a __call__ method to convert arguments to the desired type. UnitSpec might look something like (again in broken code)

class UnitSpec:
    def __init__(self, unit_or_physical_type, dtype=Quantity):
        self.unit_or_physical_type = unit_or_physical_type
        self.dtype= dtype  # allows for subclasses of Quantity, like Angle or Distance

    def _cast_unit(self, arg, unit, dtype): return dtype(arg, unit, copy=False)
    def _cast_physical_type(self, arg, physical_type, dtype): return self._cast_unit(arg, Unit._from_physical_type_id(_unit_physical_mapping[target]), dtype)

    def __call__(self, arg, dtype=None):
       if dtype is None: dtype=self.dtype
       # do conversion here
       if self.unit_or_physical_type is Unit:
            return self._cast_unit(arg, self.unit_or_physical_type, dtype=dtype)
       if self.unit_or_physical_type is str:
            return self._cast_physica_type(arg, self.unit_or_physical_type, dtype=dtype)
       if self.unit_or_physical_type is Sequence:
            # haven't looked deeply into this, but something similar to above

With this method I think it's even pretty easy to extend quantity_input to variable args and kwargs.

@quantity_input
def func(*args: Annotated[Quantity, UnitSpec(u.km)]):
    pass

Again, the nice thing about this paradigm is that all the heavy lifting is pre-computed when wrapper is created. Every time wapper is called, all it does is loop through the annotations and convert the input according to the unitspec call method. Since both the annotations and the equivalencies are stored in the wrapper, they can be changed by the user, after creation. Also, using quantity_input as a decorator by function call (quantity_input(func)) allows for the same function to be wrapped multiple different ways.

Thoughts?

Also, make a draft PR?

@nstarman

@nstarman - I think the draft makes a lot of sense, but I think it does more than the current decorator! In particular, the current decorator only checks that the input is a quantity with an equivalent unit - it does not convert anything. But I nevertheless think your implementation is very good; essentially, one could imagine different UnitSpec classes, some which only validate (like the current _validate_arg_value function), and some that also convert. Indeed, there has been a long-standing request for a decorator that takes a quantity and effectively does .to_value() on it, so that inside the code one can assume a plain ndarray in the requested units (and thus calculate faster); see #7368.

I could imaging having a keyword argument to decorator that could either be a class that handles things or a descriptive string, such as action='validate' (default), action='to', `action='to_value'.

Above said: I think much of it might also be implementable by just subclassing the QuantityInput class... Still, there does seem to be a bit of a separation of concerns between the general wrapping and the action one takes on a given input.

idea for dequantify_input decorator

I've run into this issue in my codebase and an easy fix would be to use similar methods already implemented for annotations on arguments and keyword arguments.

             valid_empty = (inspect.Signature.empty, None)
-            if wrapped_signature.return_annotation not in valid_empty:
+            if (wrapped_signature.return_annotation not in valid_empty) and isinstance(
+                wrapped_signature.return_annotation, (str, UnitBase, FunctionUnitBase)
+            ):
                 return return_.to(wrapped_signature.return_annotation)
             else:
                 return return_

If a solution like this would be acceptable I already have the patch on a branch of mine. Thought's on this?
Edit: I'm aware there is some bigger refactoring going on, but maybe in the meantime some nuisance like this can be fixed anyway.

@mhvk

in the meantime some nuisance like this can be fixed anyway.

Sounds good to me! I defer approval of a such a solution to @mhvk and the maintainers.

I think #10662 was close enough that really we should just continue with that first (and perhaps use the above suggestion if not done already) -- #10662 is hanging there for a lack of time that a new PR won't resolve...

I would just be aiming for a bugfix, in case the other PR does not land in time. Maybe it would even possible to backport this fix to existing releases. To me not being able to use non-unit return annotations with the decorator is more a bug than a missing feature. I will just post a PR and if it's not worth it you can just close it.

Sure, thanks!!

Navigation Menu

Repository navigation

Allow return annotations other than Quantity #10156

Description

Expected behavior

Actual behavior

Steps to Reproduce

System Details

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Allow return annotations other than Quantity #10156

Description

Description

Expected behavior

Actual behavior

Steps to Reproduce

System Details

Activity

Cadair commented on Apr 16, 2020

nstarman commented on Aug 8, 2020

mhvk commented on Aug 8, 2020

nstarman commented on Aug 9, 2020

nstarman commented on Aug 9, 2020

nstarman commented on Aug 10, 2020

mhvk commented on Aug 10, 2020

nstarman commented on Aug 10, 2020

mhvk commented on Aug 11, 2020

nstarman commented on Aug 13, 2020

mhvk commented on Aug 13, 2020

mhvk commented on Aug 13, 2020

noc0lour commented on Apr 8, 2021

nstarman commented on Apr 8, 2021

mhvk commented on Apr 9, 2021

noc0lour commented on Apr 9, 2021

mhvk commented on Apr 9, 2021

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions