Update dataarray_plot.py

pydata · Mar 23, 2024 · 1ac134d · 1ac134d
1 parent adc945a
commit 1ac134d
Showing 1 changed file with 51 additions and 26 deletions.
diff --git a/xarray/plot/dataarray_plot.py b/xarray/plot/dataarray_plot.py
@@ -333,7 +333,8 @@ def line(  # type: ignore[misc,unused-ignore]  # None is hashable :(
     add_legend: bool = True,
     _labels: bool = True,
     **kwargs: Any,
-) -> list[Line3D]: ...
+) -> list[Line3D]:
+    ...
 
 
 @overload
@@ -360,7 +361,8 @@ def line(
     add_legend: bool = True,
     _labels: bool = True,
     **kwargs: Any,
-) -> FacetGrid[T_DataArray]: ...
+) -> FacetGrid[T_DataArray]:
+    ...
 
 
 @overload
@@ -387,7 +389,8 @@ def line(
     add_legend: bool = True,
     _labels: bool = True,
     **kwargs: Any,
-) -> FacetGrid[T_DataArray]: ...
+) -> FacetGrid[T_DataArray]:
+    ...
 
 
 # This function signature should not change so that it can use
@@ -541,7 +544,8 @@ def step(  # type: ignore[misc,unused-ignore]  # None is hashable :(
     row: None = None,  # no wrap -> primitive
     col: None = None,  # no wrap -> primitive
     **kwargs: Any,
-) -> list[Line3D]: ...
+) -> list[Line3D]:
+    ...
 
 
 @overload
@@ -554,7 +558,8 @@ def step(
     row: Hashable,  # wrap -> FacetGrid
     col: Hashable | None = None,
     **kwargs: Any,
-) -> FacetGrid[DataArray]: ...
+) -> FacetGrid[DataArray]:
+    ...
 
 
 @overload
@@ -567,7 +572,8 @@ def step(
     row: Hashable | None = None,
     col: Hashable,  # wrap -> FacetGrid
     **kwargs: Any,
-) -> FacetGrid[DataArray]: ...
+) -> FacetGrid[DataArray]:
+    ...
 
 
 def step(
@@ -1140,7 +1146,8 @@ def scatter(  # type: ignore[misc,unused-ignore]  # None is hashable :(
     extend: ExtendOptions = None,
     levels: ArrayLike | None = None,
     **kwargs,
-) -> PathCollection: ...
+) -> PathCollection:
+    ...
 
 
 @overload
@@ -1181,7 +1188,8 @@ def scatter(
     extend: ExtendOptions = None,
     levels: ArrayLike | None = None,
     **kwargs,
-) -> FacetGrid[T_DataArray]: ...
+) -> FacetGrid[T_DataArray]:
+    ...
 
 
 @overload
@@ -1222,7 +1230,8 @@ def scatter(
     extend: ExtendOptions = None,
     levels: ArrayLike | None = None,
     **kwargs,
-) -> FacetGrid[T_DataArray]: ...
+) -> FacetGrid[T_DataArray]:
+    ...
 
 
 @_plot1d
@@ -1689,7 +1698,8 @@ def imshow(  # type: ignore[misc,unused-ignore]  # None is hashable :(
     ylim: ArrayLike | None = None,
     norm: Normalize | None = None,
     **kwargs: Any,
-) -> AxesImage: ...
+) -> AxesImage:
+    ...
 
 
 @overload
@@ -1729,7 +1739,8 @@ def imshow(
     ylim: ArrayLike | None = None,
     norm: Normalize | None = None,
     **kwargs: Any,
-) -> FacetGrid[T_DataArray]: ...
+) -> FacetGrid[T_DataArray]:
+    ...
 
 
 @overload
@@ -1769,7 +1780,8 @@ def imshow(
     ylim: ArrayLike | None = None,
     norm: Normalize | None = None,
     **kwargs: Any,
-) -> FacetGrid[T_DataArray]: ...
+) -> FacetGrid[T_DataArray]:
+    ...
 
 
 @_plot2d
@@ -1848,9 +1860,10 @@ def _center_pixels(x):
         # missing data transparent.  We therefore add an alpha channel if
         # there isn't one, and set it to transparent where data is masked.
         if z.shape[-1] == 3:
-            alpha = np.ma.ones(z.shape[:2] + (1,), dtype=z.dtype)
+            safe_dtype = np.promote_types(z.dtype, np.uint8)
+            alpha = np.ma.ones(z.shape[:2] + (1,), dtype=safe_dtype)
             if np.issubdtype(z.dtype, np.integer):
-                alpha = 255 * alpha
+                alpha[:] = 255
             z = np.ma.concatenate((z, alpha), axis=2)
         else:
             z = z.copy()
@@ -1905,7 +1918,8 @@ def contour(  # type: ignore[misc,unused-ignore]  # None is hashable :(
     ylim: ArrayLike | None = None,
     norm: Normalize | None = None,
     **kwargs: Any,
-) -> QuadContourSet: ...
+) -> QuadContourSet:
+    ...
 
 
 @overload
@@ -1945,7 +1959,8 @@ def contour(
     ylim: ArrayLike | None = None,
     norm: Normalize | None = None,
     **kwargs: Any,
-) -> FacetGrid[T_DataArray]: ...
+) -> FacetGrid[T_DataArray]:
+    ...
 
 
 @overload
@@ -1985,7 +2000,8 @@ def contour(
     ylim: ArrayLike | None = None,
     norm: Normalize | None = None,
     **kwargs: Any,
-) -> FacetGrid[T_DataArray]: ...
+) -> FacetGrid[T_DataArray]:
+    ...
 
 
 @_plot2d
@@ -2038,7 +2054,8 @@ def contourf(  # type: ignore[misc,unused-ignore]  # None is hashable :(
     ylim: ArrayLike | None = None,
     norm: Normalize | None = None,
     **kwargs: Any,
-) -> QuadContourSet: ...
+) -> QuadContourSet:
+    ...
 
 
 @overload
@@ -2078,7 +2095,8 @@ def contourf(
     ylim: ArrayLike | None = None,
     norm: Normalize | None = None,
     **kwargs: Any,
-) -> FacetGrid[T_DataArray]: ...
+) -> FacetGrid[T_DataArray]:
+    ...
 
 
 @overload
@@ -2118,7 +2136,8 @@ def contourf(
     ylim: ArrayLike | None = None,
     norm: Normalize | None = None,
     **kwargs: Any,
-) -> FacetGrid[T_DataArray]: ...
+) -> FacetGrid[T_DataArray]:
+    ...
 
 
 @_plot2d
@@ -2171,7 +2190,8 @@ def pcolormesh(  # type: ignore[misc,unused-ignore]  # None is hashable :(
     ylim: ArrayLike | None = None,
     norm: Normalize | None = None,
     **kwargs: Any,
-) -> QuadMesh: ...
+) -> QuadMesh:
+    ...
 
 
 @overload
@@ -2211,7 +2231,8 @@ def pcolormesh(
     ylim: ArrayLike | None = None,
     norm: Normalize | None = None,
     **kwargs: Any,
-) -> FacetGrid[T_DataArray]: ...
+) -> FacetGrid[T_DataArray]:
+    ...
 
 
 @overload
@@ -2251,7 +2272,8 @@ def pcolormesh(
     ylim: ArrayLike | None = None,
     norm: Normalize | None = None,
     **kwargs: Any,
-) -> FacetGrid[T_DataArray]: ...
+) -> FacetGrid[T_DataArray]:
+    ...
 
 
 @_plot2d
@@ -2355,7 +2377,8 @@ def surface(
     ylim: ArrayLike | None = None,
     norm: Normalize | None = None,
     **kwargs: Any,
-) -> Poly3DCollection: ...
+) -> Poly3DCollection:
+    ...
 
 
 @overload
@@ -2395,7 +2418,8 @@ def surface(
     ylim: ArrayLike | None = None,
     norm: Normalize | None = None,
     **kwargs: Any,
-) -> FacetGrid[T_DataArray]: ...
+) -> FacetGrid[T_DataArray]:
+    ...
 
 
 @overload
@@ -2435,7 +2459,8 @@ def surface(
     ylim: ArrayLike | None = None,
     norm: Normalize | None = None,
     **kwargs: Any,
-) -> FacetGrid[T_DataArray]: ...
+) -> FacetGrid[T_DataArray]:
+    ...
 
 
 @_plot2d