Merge pull request matplotlib#28755 from QuLogic/tri-typing

Author: timhoffm

lib/matplotlib/_tri.pyi

@@ -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: ...

lib/matplotlib/tests/test_triangulation.py

@@ -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)
 
@@ -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
@@ -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(
@@ -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():

src/tri/_tri.cpp

@@ -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");