Skip to content

TYP: Add typing for internal _tri extension #28755

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Sep 7, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 26 additions & 16 deletions lib/matplotlib/_tri.pyi
Original file line number Diff line number Diff line change
@@ -1,26 +1,36 @@
# This is a private module implemented in C++
# As such these type stubs are overly generic, but here to allow these types
# as return types for public methods
from typing import Any, final
from typing import final

import numpy as np
import numpy.typing as npt

@final
class TrapezoidMapTriFinder:
def __init__(self, *args, **kwargs) -> None: ...
def find_many(self, *args, **kwargs) -> Any: ...
def get_tree_stats(self, *args, **kwargs) -> Any: ...
def initialize(self, *args, **kwargs) -> Any: ...
def print_tree(self, *args, **kwargs) -> Any: ...
def __init__(self, triangulation: Triangulation): ...
def find_many(self, x: npt.NDArray[np.float64], y: npt.NDArray[np.float64]) -> npt.NDArray[np.int_]: ...
def get_tree_stats(self) -> list[int | float]: ...
def initialize(self) -> None: ...
def print_tree(self) -> None: ...

@final
class TriContourGenerator:
def __init__(self, *args, **kwargs) -> None: ...
def create_contour(self, *args, **kwargs) -> Any: ...
def create_filled_contour(self, *args, **kwargs) -> Any: ...
def __init__(self, triangulation: Triangulation, z: npt.NDArray[np.float64]): ...
def create_contour(self, level: float) -> tuple[list[float], list[int]]: ...
def create_filled_contour(self, lower_level: float, upper_level: float) -> tuple[list[float], list[int]]: ...

@final
class Triangulation:
def __init__(self, *args, **kwargs) -> None: ...
def calculate_plane_coefficients(self, *args, **kwargs) -> Any: ...
def get_edges(self, *args, **kwargs) -> Any: ...
def get_neighbors(self, *args, **kwargs) -> Any: ...
def set_mask(self, *args, **kwargs) -> Any: ...
def __init__(
self,
x: npt.NDArray[np.float64],
y: npt.NDArray[np.float64],
triangles: npt.NDArray[np.int_],
mask: npt.NDArray[np.bool_] | tuple[()],
edges: npt.NDArray[np.int_] | tuple[()],
neighbors: npt.NDArray[np.int_] | tuple[()],
correct_triangle_orientation: bool,
): ...
def calculate_plane_coefficients(self, z: npt.ArrayLike) -> npt.NDArray[np.float64]: ...
def get_edges(self) -> npt.NDArray[np.int_]: ...
def get_neighbors(self) -> npt.NDArray[np.int_]: ...
def set_mask(self, mask: npt.NDArray[np.bool_] | tuple[()]) -> None: ...
38 changes: 19 additions & 19 deletions lib/matplotlib/tests/test_triangulation.py
Original file line number Diff line number Diff line change
Expand Up @@ -1181,43 +1181,44 @@ def test_tricontourf_decreasing_levels():
plt.tricontourf(x, y, z, [1.0, 0.0])


def test_internal_cpp_api():
def test_internal_cpp_api() -> None:
# Following github issue 8197.
from matplotlib import _tri # noqa: F401, ensure lazy-loaded module *is* loaded.

# C++ Triangulation.
with pytest.raises(
TypeError,
match=r'__init__\(\): incompatible constructor arguments.'):
mpl._tri.Triangulation()
mpl._tri.Triangulation() # type: ignore[call-arg]

with pytest.raises(
ValueError, match=r'x and y must be 1D arrays of the same length'):
mpl._tri.Triangulation([], [1], [[]], (), (), (), False)
mpl._tri.Triangulation(np.array([]), np.array([1]), np.array([[]]), (), (), (),
False)

x = [0, 1, 1]
y = [0, 0, 1]
x = np.array([0, 1, 1], dtype=np.float64)
y = np.array([0, 0, 1], dtype=np.float64)
with pytest.raises(
ValueError,
match=r'triangles must be a 2D array of shape \(\?,3\)'):
mpl._tri.Triangulation(x, y, [[0, 1]], (), (), (), False)
mpl._tri.Triangulation(x, y, np.array([[0, 1]]), (), (), (), False)

tris = [[0, 1, 2]]
tris = np.array([[0, 1, 2]], dtype=np.int_)
with pytest.raises(
ValueError,
match=r'mask must be a 1D array with the same length as the '
r'triangles array'):
mpl._tri.Triangulation(x, y, tris, [0, 1], (), (), False)
mpl._tri.Triangulation(x, y, tris, np.array([0, 1]), (), (), False)

with pytest.raises(
ValueError, match=r'edges must be a 2D array with shape \(\?,2\)'):
mpl._tri.Triangulation(x, y, tris, (), [[1]], (), False)
mpl._tri.Triangulation(x, y, tris, (), np.array([[1]]), (), False)

with pytest.raises(
ValueError,
match=r'neighbors must be a 2D array with the same shape as the '
r'triangles array'):
mpl._tri.Triangulation(x, y, tris, (), (), [[-1]], False)
mpl._tri.Triangulation(x, y, tris, (), (), np.array([[-1]]), False)

triang = mpl._tri.Triangulation(x, y, tris, (), (), (), False)

Expand All @@ -1232,9 +1233,9 @@ def test_internal_cpp_api():
ValueError,
match=r'mask must be a 1D array with the same length as the '
r'triangles array'):
triang.set_mask(mask)
triang.set_mask(mask) # type: ignore[arg-type]

triang.set_mask([True])
triang.set_mask(np.array([True]))
assert_array_equal(triang.get_edges(), np.empty((0, 2)))

triang.set_mask(()) # Equivalent to Python Triangulation mask=None
Expand All @@ -1244,15 +1245,14 @@ def test_internal_cpp_api():
with pytest.raises(
TypeError,
match=r'__init__\(\): incompatible constructor arguments.'):
mpl._tri.TriContourGenerator()
mpl._tri.TriContourGenerator() # type: ignore[call-arg]

with pytest.raises(
ValueError,
match=r'z must be a 1D array with the same length as the x and y '
r'arrays'):
mpl._tri.TriContourGenerator(triang, [1])
match=r'z must be a 1D array with the same length as the x and y arrays'):
mpl._tri.TriContourGenerator(triang, np.array([1]))

z = [0, 1, 2]
z = np.array([0, 1, 2])
tcg = mpl._tri.TriContourGenerator(triang, z)

with pytest.raises(
Expand All @@ -1263,13 +1263,13 @@ def test_internal_cpp_api():
with pytest.raises(
TypeError,
match=r'__init__\(\): incompatible constructor arguments.'):
mpl._tri.TrapezoidMapTriFinder()
mpl._tri.TrapezoidMapTriFinder() # type: ignore[call-arg]

trifinder = mpl._tri.TrapezoidMapTriFinder(triang)

with pytest.raises(
ValueError, match=r'x and y must be array-like with same shape'):
trifinder.find_many([0], [0, 1])
trifinder.find_many(np.array([0]), np.array([0, 1]))


def test_qhull_large_offset():
Expand Down
2 changes: 1 addition & 1 deletion src/tri/_tri.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1314,7 +1314,7 @@ TrapezoidMapTriFinder::TriIndexArray
TrapezoidMapTriFinder::find_many(const CoordinateArray& x,
const CoordinateArray& y)
{
if (x.ndim() != 1 || x.shape(0) != y.shape(0))
if (x.ndim() != 1 || y.ndim() != 1 || x.shape(0) != y.shape(0))
throw std::invalid_argument(
"x and y must be array-like with same shape");

Expand Down
Loading