Typing to help call functions with the correct arguments

[DocOtak]

I use vscode a lot and like how helpful type hints can be, I threw together a little experiment to see if it could help with having the correct arguments to functions that need certain specific outputs of other functions (in most oceanographic contexts), I'm attaching it here but this doesn't need to go anywhere as a proposal. I didn't check or even try to see if this could be done in the generic wrapped way we have done everything else. If the functions need to actually exist for things like pylance or mypy to be useful... well I hope it doesn't need that.

from typing import NewType, cast

from numpy.typing import ArrayLike

SA = NewType("SA", object)
SP = NewType("SP", object)
SR = NewType("SR", object)
CT = NewType("CT", object)
P = NewType("P", object)
Rho = NewType("Rho", object)


def SA_from_SP(SP: SP | ArrayLike, p: P | ArrayLike, lon, lat) -> SA:
    from gsw_xarray import SA_from_SP as _func
    return _func(SP, p, lon, lat)

def SP_from_C(C, t, p: P | ArrayLike) -> SP:
    from gsw_xarray import SP_from_C as _func
    return _func(C, t, p)

def SR_from_SP(SP: SP| ArrayLike) -> SR:
    from gsw_xarray import SR_from_SP as _func
    return _func(SP)

def rho(SA: SA| ArrayLike, CT: CT| ArrayLike, p: P | ArrayLike) -> Rho:
    from gsw_xarray import rho as _func
    return _func(SA, CT, p)

sal = SP_from_C(0, 0, 0)

sa = SA_from_SP(0, 0, 0, 0)

# type check fails as it should (at least with pylance in vscode)
SA_from_SP(SA_from_SP(0,0,0,0), 0,0,0)
#          ^^^^^^^^^^^^^^^^^^^ - vscode red wiggles this part with the message
# Argument of type "SA" cannot be assigned to parameter "SP" of type "ArrayLike | SP" in function "SA_from_SP"
#   Type "SA" is not assignable to type "ArrayLike | SP"
#     "SA" is not assignable to "SP"
#     "SA" is incompatible with protocol "_Buffer"
#       "__buffer__" is not present
#     "SA" is incompatible with protocol "_SupportsArray[dtype[Any]]"
#       "__array__" is not present
#     "SA" is incompatible with protocol "_NestedSequence[_SupportsArray[dtype[Any]]]"
#       "__len__" is not present
#
# Other wiggles marked but I'm not including the message


# what do I do if no lon/lat to to the SA conversion
# can use SR instead if ok with the extra error/uncertainty
# make you cast it so you _know_ you are being naughty
sr = SR_from_SP(35)

# type check error
rho(sr, 0,0)
#   ^^  - red wiggles

# this is OK
rho(cast(SA, sr), 0,0)