I’d like to add typehints to a function that accepts either np.float32 arrays or np.float64 arrays and returns the same type:

from typing import overload, Union

import numpy as np
import numpy.typing as npt

NPArray_FLOAT32 = npt.NDArray[np.float32]
NPArray_FLOAT64 = npt.NDArray[np.float64]
NPArray_FLOAT32_64 = Union[NPArray_FLOAT32, NPArray_FLOAT64]

@overload
def foo(xa: NPArray_FLOAT32, xb: NPArray_FLOAT32) -> NPArray_FLOAT32: ...


@overload
def foo(xa: NPArray_FLOAT64, xb: NPArray_FLOAT64) -> NPArray_FLOAT64: ...


def foo(xa: NPArray_FLOAT32_64, xb: NPArray_FLOAT32_64) -> NPArray_FLOAT32_64:
    # ...

However, this results in the following error from mypy

mypy [overload-overlap]: Overloaded function signatures 1 and 2 overlap with incompatible return types.

What’s the right way to do this? This almost seems like a bug with mypy since np.float32 does not overlap with np.float64.