Annotate a tuple containing an NDArray

Question

I am trying to annotate a function that returns a tuple: tuple[NDArray[Any, Int32 | Float32], dict[str, Any]]: using nptyping for NDArray and typing for other python objects. Here is some code:

from nptyping import NDArray, UInt8, Int32, Float32, Shape, Bool
import rasterio
from typing import Generator, Any, Final, Literal

    
    
    def _scale_and_round(
        self, arr: NDArray[Any, Float32]
    ) -> tuple[NDArray[Any, Int32 | Float32], dict[str, Any]]:
        
        array: NDArray[Any, Any] = arr * self.scale_factor
        if self.scale_factor == 1000:
            array = array.astype(np.int32)
        return array, self.metadata

    def ndvi(
        self, red_src: Any, nir_src: Any
    ) -> tuple[NDArray[Any, Int32 | Float32], dict[str, Any]]:
        
        redB: NDArray[Any, Any] = red_src.read()
        nirB: NDArray[Any, Any] = nir_src.read()
        np.seterr(divide="ignore", invalid="ignore")
        ndvi: NDArray[Any, Float32] = (
            nirB.astype(np.float32) - redB.astype(np.float32)
        ) / (nirB.astype(np.float32) + redB.astype(np.float32))
        # replace nan with 0
        where_are_NaNs: NDArray[Any, Bool] = np.isnan(ndvi)
        ndvi[where_are_NaNs] = 0

        return self._scale_and_round(ndvi)

When trying to run this code I get:

nptyping.error.InvalidArgumentsError: Unexpected argument of type <class 'type'>, expecting a string.

Pointing to the line -> -> tuple[NDArray[Any, Int32 | Float32], dict[str, Any]]:

Does somebody know how Am I suppose to annotate ndvi func? Thanks

Answer 1

Separate the type: Int32 | Float32. So instead of this:

tuple[NDArray[Any, Int32 | Float32], dict[str, Any]]

you'll have this:

tuple[NDArray[Any, Int32] | NDArray[Any, Float32], dict[str, Any]]

Numpy ndarrays are homogeneous, it means that the items in the array have to be of the same type. First syntax would suggest that the array could contain items of both types at the same time - which is not true. In this case, it can be either array of integers or floats.