Description

This is a big set of changes, PLEASE HELP.

This DRAFT pull request is to coordinate a number of PRs improving quantity_input and unit-aware static typing:

The goals are:

1. Use quantity_input with statically typed code
2. Allow SpecificTypeQuantity subclasses, like Distance or Angle, to be used instead of Quantity in the input verification of quantity_input
3. Speed up calling quantity_input-wrapped functions by pre-computing as much of the conversion mechanics at wrapper initialization.
4. a) Allow inputs to be cast to specific units (or physical types), not just verified
   b) post-cast, to pass input to function as .to_value() so that functions can work with ndarrays guaranteed to be in the correct units
5. Expose all of these improvements to the user in a python-standard way, so that they can be changed even after the function has been initialized.

Select Ancillary Benefits:
See whole comment for details

• unit-aware static type hints can be used independently of quantity_input
• UnitSpec classes can be used independently of quantity_input to provide specific unit/physical-type validation/conversion/etc.
• User-defined UnitSpec subclasses are automatically registered and usable by UnitSpec and quantity_input

Summary of each PR

1. add unit-aware quantity annotations. #10662 : Introduce unit aware annotations.

• set Annotated[Quantity, <unit>] as official unit-aware Quantity annotation
• add __class_getitem__ method to Quantity to shortcut creating.
  Caveat: py3.7+
• create stand-in Annotated class. Not valid for static typing. Run-time only. But same API as static implementation.
• work with quantity_input

2. Rewrite quantity_input internals so that all information is stored in wrapper.__annotations__, converting valid inputs to Quantity[<unit or pt>] from PR#1. This is still not valid static typing, but it will trivially solve Allow return annotations other than Quantity #10156 by having a) non-unit / physicaltype return annotations ignored and b) allowing quantity_input(return=None) to store None in wrapper.__annotations__ instead of Quantity[<unit or pt>]

3. UnitSpec to solve Decorator for converting quantity to arrays in specific units? #7368 and multiple annotations

• Extend Unit to your UnitSpec so that quantity_input becomes more powerful (addressing Decorator for converting quantity to arrays in specific units? #7368)
• Change quantity_input so that it uses the actual class, making Angle and Distance possible (maybe will be easiest if done together with (2)).
  have quantity_input convert this to Annotated[Quantity, UnitSpec(<unit>), other] for internal use.
  Add method to Quantity to create UnitSpec annotations

4. switch to valid static typing, having had time to discuss this more broadly with the rest of Astropy.

Overall Summary

Subject to change.

Static Typing

Current State:

There is no means to create unit-aware type annotations. For instance, the annotations used by quantity_input, like u.km or distance, are not compliant with static typing since the inputs are of (sub)type Quantity.

Proposed Change:

This problem can be solved by the typing.Annotated class, from this PEP which allows valid static type annotations to include metadata. Now, unit-aware type annotations will be of the form Annotated[Quantity, <unit>, other_annotations], where the unit must be the first annotation.

A __class_getitem__ method will be added to Quantity to shortcut creating unit aware annotations.
Quantity[<unit>, other_annotations] == Annotated[Quantity, <unit>, other_annotation].
As a class method, this will similarly work for all subclasses of Quantity.

quantity_input will accept unit-aware quantity annotations of this type. Both the type (eg Quantity) and unit must be valid inputs.

Implementation Details:

typing_extensions is the official backport for Annotated and works on any py3.5+. Until py3.9+ is the minimum python version, when Annotated is in the built-in typing library, we use the typing_extension backport. Until the typing_extension backport is added as an astropy dependency we emulate API for unsupported features.

Pre-Computation and User Accessibility

Current State:

quantity_input uses function annotations to determine the required argument units and since it also uses functools.wraps those annotations are propagated to the wrapper returned by quantity_input. However, other inputs to quantity_input, such as keyword arguments are not similarly stored.

Proposed Change:

The wrapper's __annotations__ will store all unit/physical_type information, merging the information provided from the keyword arguments with the original function's __annotations__. The original function will be unchanged and accessible in the __wrapped__ attribute. Since decorators can be applied as functions, not just with the "pie" syntax, this will allow multiple different versions of the same function to be created with different unit specifications. Not only that, but storing all information in __annotations__ means that the desired units are pre-computed, increasing function call speed, and are also editable by knowledgeable users in a python-standard format.

UnitSpec

For quantity_input, the Annotated[Quantity, <unit>, other_annotations] is not specific enough, since it only annotates which argument should be processed, not what to actually do with them. We introduce a UnitSpec (abbreviation for "unit specification") class to do both.

quantity_input will accept the standard Annotated[Quantity, <unit>, other_annotations], but also accept the more detailed annotation Annotated[Quantity, UnitSpec(<unit>), other_annotations]. Internally, any of non-unitspec annotations will be converted.

The UnitSpec.__call__ will control how quantity_input interacts with each input — if it just verifies unit / physical_type, or casts to unit / physical_type, or does to_value, etc. UnitSpec will enable different treatment of each input: for a function func(x, y), x might just be verified, but y = y.to_value(u.km).

UnitSpec will have the following pre-specified options for the __call__ output, set with the "action" attribute.

• "validate" : to only validate physical type
• "convert" : to convert quantities to specific units / physical type
• "to_value" : takes .to_value after conversion
• "from_value" : reassigns units if none present & verifies units if present.

The following subclasses of UnitSpec will also be implemented:

• UnitSpecValidate : which can only validate physical type
• UnitSpecConvert : which will convert quantities to specific units / physical type
• UnitSpecValue : a subclass of UnitSpecConvert which takes .to_value after conversion
• UnitSpecAssign : which reassigns units if none present & verifies units if present.

Using __init_subclass__ these and any user-created subclasses of UnitSpec will be added to a registry. The "action" kwarg of UnitSpec must be a key in this registry and will be used by UnitSpec to set the behavior of UnitSpec.__call__. In this way UnitSpec can stand in for any of its subclasses, which are only really needed to make a behavior immutable.

UnitSpec will drastically simplify a few things in quantity_input, and potentially elsewhere.

For instance, in quantity_input, all that's needed to check a multiply-annotated argument and then perform the validation is

if isAnnotated(antn):
    md = [isinstance(x, UnitSpecBase) for x in antn.__metadata__]
    if not any(md):  # skip to next argument
        continue
    elif sum(md) > 1 :  # more than one UnitSpec
        raise Exception()

    unitspec = antn.__metadata__[np.nonzero(md)]
    if name == "return":  # save the return casting until actually call function.
        return_spec = unitspec
    else:
        ba.arguments[name] = unitspec(ba.arguments[name])

quantity_input implementation

Order of precedence of unit specifications:

1. Kwargs into quantity_input:
   quantity_input(x=u.hour)(foo).
   This allows quantity_input to wrap functions from external libraries without modifying the function.
2. function annotations
   So that calling quantity_input as a function overrides the default unit specifications.
3. units from argument default values
   To respect type hints

Argument-specific behavior in quantity_input is allowed by mixing UnitSpec subclasses in wrapper.__annotations__.

For simple means of ensuring uniform UnitSpec behavior, there will be a few subclasses of QuantityInput which convert any user-created unit-aware Quantity annotations to the appropriate UnitSpec subclasses. The complexities of the wrapper construction can happen in super and only ever need to be written once. The UnitSpec auto-registry very easily allows user-created UnitSpec and therefore user-created QuantityInput subclasses.

• ValidateQuantityInput : converts UnitSpec to UnitSpecValidate
• ConvertQuantityInput : converts UnitSpec to UnitSpecConvert
• DeQuantityInput : converts UnitSpec to UnitSpecValue

For example, to make a to_value flavor of quantity_input is a simple for-loop replacing UnitSpec with UnitSpecValue, the subclass that performs conversion and takes the value. While this is the default behavior for UnitSpec and the subclass used by UnitSpec(...action = 'value'), the subclass is immutable and more explicit.

class DeQuantityInput(QuantityInput):
    """QuantityInput, with ``to_value(unit)`` on all specified inputs."""

    def __call__(self, wrapped_function):

        # use all the machinery from superclass
        wrapper = super().__call__(wrapped_function)

        # but now need to ensure that everything is correct UnitSpec
        _antns_ = dict()
        for name, antn in wrapper.__annotations__.items():

            if isAnnotated(antn):
                unitspec = antn.__metadata__[0]
                if isinstance(unitspec, UnitSpecBase):
                    if name == "return":
                        _antns_[name] = UnitSpec(unitspec, action="from_value")
                    else:
                        _antns_[name] = UnitSpecValue(unitspec)

        wrapper.__annotations__.update(_antns_)

        return wrapper

Examples

Assuming the following imports

import typing_extensions as T
from typing_extensions import Annotated
import astropy.units as u
from astropy.units import quantity_input, Quantity, UnitSpec
import astropy.coordinates as coord

Example 1)

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

# annotations
foo.__annotations__

{
"x": Annotated[Quantity, UnitSpec(u.m | validate Quantity)],
"v": Annotated[Quantity, UnitSpec("speed" | validate Quantity)],
"t": Annotated[Quantity, UnitSpec(u.s | validate Quantity)],
"y": T.Union[Annotated[Quantity, UnitSpec(u.km | validate Quantity)], NoneType]
"flag": T.Union[str, NoneType],
"return": Annotated[Quantity, UnitSpec("distance" | validate Quantity)]
}

# and the original function
foo.__wrapped__.__annotations__

{
"x": u.m, "v": Quantity, "flag": T.Optional[str],
"y": T.Union[Annotated[Quantity, u.km], NoneType], "return": "distance"
}

To force x input to be Distance or subclass

foo.__annotations__["x"].__origin__ = coord.Distance
foo.__annotations__["x"].__metadata__[0].dtype = coord.Distance
foo.__annotations__["x"]

Annotated[Distance, UnitSpec(u.m | validate Distance)]

To change t to action "convert"

foo.__annotations__["t"].__metadata_[0].action = "convert"
foo.__annotations__["t"]

Annotated[Quantity, UnitSpec(u.s | convert Quantity)]

To make a to_value flavor of quantity_input is a simple for-loop replacing UnitSpec with UnitSpecValue, the subclass that performs conversion and takes the value. While this is the default behavior for UnitSpec and the subclass used by UnitSpec(...action = 'value'), the subclass is immutable and more explicit.

class DeQuantityInput(QuantityInput):
    """QuantityInput, with ``to_value(unit)`` on all specified inputs."""

    def __call__(self, wrapped_function):

        # use all the machinery from superclass
        wrapper = super().__call__(wrapped_function)

        # but now need to ensure that everything is correct UnitSpec
        _antns_ = dict()
        for name, antn in wrapper.__annotations__.items():

            if isAnnotated(antn):
                unitspec = antn.__metadata__[0]
                if isinstance(unitspec, UnitSpecBase):
                    if name == "return":
                        _antns_[name] = UnitSpec(unitspec, action="from_value")
                    else:
                        _antns_[name] = UnitSpecValue(unitspec)

        wrapper.__annotations__.update(_antns_)

        return wrapper

Fixes #10156, #5606, #7368