How to create custom input types for function defenition for external libraries in Python?

Question

I wanted to write functions with detailed types so that during usage of any typecheker e.g. Pylance in VScode you could clearly see the input type. I managed to achieve some result with default input types like int and then extended it to numpy.ndarray. But this is where I ran into an issue. The code for example function is below:

import numpy as np

def get_distance(arr1: np.ndarray, arr2: np.ndarray) -> int:
    """Returns Euclidian distance between 2 points in 3D space

    Args:
        arr1 (np.ndarray[float]): array 1 format [[x,y,z],...]
        arr2 (np.ndarray[float]): array 2 format [[x,y,z],...]

    Returns:
        [int]: Distance between 2 points
    """
    ...

return some_distance

When I defined it and then looked for definition under Pylance I managed to get np.ndarray[Unknown, Unknown]:

(function) get_distance: (arr1: ndarray[Unknown, Unknown], arr2: ndarray[Unknown, Unknown]) -> int 

Returns Euclidian distance between 2 points in 3D space

Args:
    arr1 (np.ndarray[float]): point array 1
    arr2 (np.ndarray[float]): point array 2

Returns:
    [int]: Distance between 2 points

My question is: Is there any way to define an input type such that even when using numpy or any other external library the correct input type would be displayed? For example np.ndarray[float] or something similar.

Answer 1

is this what you want?

from numpy.typing import NDArray
import numpy as np

def get_distance(
  arr1: NDArray[np.float64], 
  arr2: NDArray[np.float64]
) -> int:

    # ...

    return some_distance

note that you need to use the typing submodule from numpy.

for me this method is working.

but note again that this method is not working: arr: np.ndarray[float]

as you can see in my editor with LSP-pyright

this is underlined with red (error)

but this method is very fine:

note again again: you can use any float you like

def get_distance(arr1: NDArray[np.float32] ...
def get_distance(arr1: NDArray[np.float64] ...
def get_distance(arr1: NDArray[np.float128] ...
# also integer
def get_distance(arr1: NDArray[np.int32],

but, most important, the data type you use should be defined in the numpy module and you have to use typing from numpy, not from standard python. why? because numpy has its own datatypes similar to C, because numpy is written in C and compiled and imported in python.

Answer 2

Your problem is specific to ndarray -> its typing has 2 arguments - shape and data type. TypeError: 'type' object is not subscriptable is because you are passing a float as a shape argument, which is supposed to be list-like. (Although personally I got more clear error on Python 3.10 - TypeError: Too few arguments for numpy.ndarray

This will work:

from typing import Any

def get_distance(arr1: np.ndarray[Any, float], arr2: np.ndarray[Any, float]) -> int: