TYP: Add typing for internal _tri extension

QuLogic · QuLogic · commit 414b63949aa9 · 2024-08-23T02:17:30.000-04:00
diff --git a/lib/matplotlib/_tri.pyi b/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) -> None: ...
+    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]) -> None: ...
+    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_],
+        edges: npt.NDArray[np.int_],
+        neighbors: npt.NDArray[np.int_],
+        correct_triangle_orientation: bool,
+    ) -> None: ...
+    def calculate_plane_coefficients(self, z: npt.NDArray[np.float64]) -> 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_]) -> None: ...
