Create standard units / quantity types

[tjgalvin]

@AlecThomson has found that in some cases astropy unit quantities are not properly serialised and deserialised when running in a dask environment. This leads to strange issues described here: https://github.com/astropy/astropy/issues/11317

It is a little confusing -- sometimes it works, and sometimes it does not. Likely some race condition. Who knows.

So long as an astropy unit quantity class is not exchanged between dask @delayed or prefect @task decorated functions there is no issue. But for cases where it is reasonable that these quantities are exchanged, it is probably safer to define a set of classes / types that clearly highlight the intent that a vartiable is of a particular unit. Especially in a way where the type checker can be leveraged.

I will add some examples to how this might be implemented.

[tjgalvin]

Something I am playing around with.

import astropy.units as u

class Radian(float):
    """Placeholder for a radian"""

    def __repr__(self) -> str:
        return f"{float(self)} rad"

def some_example(a: Radian) -> Radian:
    return Radian(a * 2)

ra_rad = Radian(12)
print(ra_rad)
print(f"{isinstance(ra_rad, float)=}")
print(f"{isinstance(ra_rad, Radian)=}")

scale = ra_rad * 2
print(f"{scale=} {isinstance(scale, Radian)=}")

scale_rad = Radian(ra_rad * 2)
print(f"{scale_rad=} {isinstance(scale_rad, Radian)=}")

ra_unit = ra_rad * u.rad
print(f"ra_unit={ra_unit}")

# by overloading the string representation we can convert to a 
# astropy.units.rad by calling the u.Quantity 
ra_quantity = u.Quantity(str(ra_rad))
print(f"ra_quantity={ra_quantity}")

# Pyright will correctly report no issue here
b = some_example(a=ra_rad)

try:
    # Pyright correctly reports an issue here
    b = some_example(a=ra_quantity) # This is an error in pyright
except TypeError:
    print("This is bad")

produces:

12.0 rad
isinstance(ra_rad, float)=True
isinstance(ra_rad, Radian)=True
isinstance(scale, Radian)=False
ra_unit=12.0 rad
ra_quantity=12.0 rad
This is bad

Running pyright does correctly report the misuse of the the input ra_quantity value into the some_example function.

Any thoughts or alternatives?