Python Type Hinting for NumPy Arrays

Proper type hints prevent shape and dtype mismatches when working with NumPy arrays.

This tutorial shows you how to use type hints with NumPy arrays.

 

 

Basic Type Hinting for NumPy Arrays

You can use numpy.ndarray to represent a general NumPy array in your type hints.

from typing import Any
import numpy as np
def calculate_sum(array: np.ndarray) -> float:
    return np.sum(array)
data = np.array([10.5, 20.3, 30.2])
result = calculate_sum(data)
print(result)

Output:

61.0

This code calculates the sum of the array [10.5, 20.3, 30.2] using np.sum and returns the result as a floating-point number.

 

Type Hinting NumPy Array Dtypes

You can specify the data type of a NumPy array using numpy.dtype.

from numpy.typing import NDArray
def filter_positive_integers(array: NDArray[np.int32]) -> NDArray[np.int32]:
    return array[array > 0]
data = np.array([-3, 2, 15, -8, 7], dtype=np.int32)
result = filter_positive_integers(data)
print(result)

Output:

[ 2 15 7]

The function filters out negative integers from the array, leaving only positive integers.

You can also use common aliases like np.float64 for specifying data types.

def calculate_mean(array: NDArray[np.float64]) -> float:
    return np.mean(array)
data = np.array([1.5, 2.7, 3.9], dtype=np.float64)
result = calculate_mean(data)
print(result)

Output:

2.6999999999999997

The function calculates the mean of the array [1.5, 2.7, 3.9] with the dtype np.float64.

If you don’t know the dtype, use Any.

def get_first_element(array: NDArray[Any]) -> Any:
    return array[0]
data = np.array(['Ahmed', 'Salma', 'Youssef'])
result = get_first_element(data)
print(result)

Output:

Ahmed

The function retrieves the first element of the array, which in this case is the string 'Ahmed'.

 

Type Hinting Array Shapes

You can specify shapes as tuples when you know the exact dimensions.

def create_identity_matrix(size: int) -> NDArray[np.float64]:
    return np.eye(size, dtype=np.float64)
result = create_identity_matrix(3)
print(result)

Output:

[[1. 0. 0.]
 [0. 1. 0.]
 [0. 0. 1.]]

The function generates a 3×3 identity matrix where diagonal elements are 1 and others are 0.

Also, you can use Ellipsis when the shape can vary but you want to hint at the dimensionality.

def normalize_data(data: NDArray[np.float64]) -> NDArray[np.float64]:
    return (data - np.mean(data)) / np.std(data)
data = np.array([[1.0, 2.0], [3.0, 4.0]])
result = normalize_data(data)
print(result)

Output:

[[-1.34164079 -0.4472136 ]
 [ 0.4472136   1.34164079]]

The function normalizes the array so that it has a mean of 0 and a standard deviation of 1.

You can combine shape and dtype for more precise type hints.

def create_matrix(rows: int, cols: int) -> NDArray[np.float64]:
    return np.random.rand(rows, cols)
result = create_matrix(2, 3)
print(result)

Output:

[[0.16321843 0.31780967 0.36052351]
 [0.0724457  0.85515167 0.00897343]]

The function creates a 2×3 matrix with random values between 0 and 1.

 

Type Hinting Array Size

Use fixed sizes to ensure you get arrays of specific sizes.

def fixed_size_array() -> NDArray[np.int_]:
    return np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
result = fixed_size_array()
print(result)

Output:

[[1 2 3]
 [4 5 6]
 [7 8 9]]

The function returns a 3×3 fixed-size array.

If you want to indicate the shape of the array in the type hint, you could use a different way with typing.Annotated:

from typing import Literal
from typing import Annotated
def fixed_size_array() -> Annotated[NDArray[np.int_], Literal[(3, 3)]]:
    return np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
result = fixed_size_array()
print(result)

Output:

[[1 2 3]
[4 5 6]
[7 8 9]]

 

Using typing.Union for multiple possible types

When an array can have multiple types, use typing.Union.

from typing import Union
def process_data(data: Union[NDArray[np.float64], NDArray[np.int32]]) -> float:
    return np.sum(data)
data_float = np.array([1.5, 2.5, 3.5], dtype=np.float64)
data_int = np.array([1, 2, 3], dtype=np.int32)
result_float = process_data(data_float)
result_int = process_data(data_int)
print(result_float)
print(result_int)

Output:

7.5
6

The function calculates the sum of the elements for both float and integer arrays.