TYP: Update typing for Mesh-like collections

greglucas · greglucas · commit d017defc5e20 · 2023-04-23T14:14:17.000-06:00
diff --git a/lib/matplotlib/collections.pyi b/lib/matplotlib/collections.pyi
@@ -205,7 +205,19 @@ class TriMesh(Collection):
     @staticmethod
     def convert_mesh_to_paths(tri: Triangulation) -> list[Path]: ...
 
-class QuadMesh(Collection):
+class _MeshData:
+    def __init__(
+        self,
+        coordinates: ArrayLike,
+        *,
+        shading: Literal["flat", "gouraud"] = ...,
+    ) -> None: ...
+    def set_array(self, A: ArrayLike | None) -> None: ...
+    def get_coordinates(self) -> ArrayLike: ...
+    def get_facecolor(self) -> ColorType | Sequence[ColorType]: ...
+    def get_edgecolor(self) -> ColorType | Sequence[ColorType]: ...
+
+class QuadMesh(_MeshData, Collection):
     def __init__(
         self,
         coordinates: ArrayLike,
@@ -217,7 +229,14 @@ class QuadMesh(Collection):
     def get_paths(self) -> list[Path]: ...
     # Parent class has an argument, perhaps add a noop arg?
     def set_paths(self) -> None: ...  # type: ignore[override]
-    def set_array(self, A: ArrayLike | None) -> None: ...
     def get_datalim(self, transData: transforms.Transform) -> transforms.Bbox: ...
-    def get_coordinates(self) -> ArrayLike: ...
     def get_cursor_data(self, event: MouseEvent) -> float: ...
+
+class PolyQuadMesh(_MeshData, PolyCollection):
+    def __init__(
+        self,
+        coordinates: ArrayLike,
+        **kwargs
+    ) -> None: ...
+    def get_facecolor(self) -> ColorType | Sequence[ColorType]: ...
+    def get_edgecolor(self) -> ColorType | Sequence[ColorType]: ...
