pydantic_numpy_array_shape_type_validation_with_numpydantic.py

1import numpy as np
2from pydantic import BaseModel
3from numpydantic import NDArray, Shape
4
5class MyModel(BaseModel):
6    # An array of any shape with float64 data
7    any_shape: NDArray[np.float64]
8    # A 2D array with specific dimensions (3 rows, 2 columns)
9    fixed_shape: NDArray[Shape["3, 2"], np.int32]
10    # A 3D array where the first two dimensions can be anything
11    wildcard_shape: NDArray[Shape["*, *, 3"], np.uint8]
12
13# Valid data
14data = {
15    "any_shape": np.array([1.0, 2.0], dtype=np.float64),
16    "fixed_shape": np.zeros((3, 2), dtype=np.int32),
17    "wildcard_shape": np.ones((10, 20, 3), dtype=np.uint8)
18}
19
20model = MyModel(**data)
21print(model.fixed_shape.shape)  # (3, 2)