Typing fixes

Author: dcherian

xarray/core/groupby.py

@@ -191,8 +191,8 @@ def values(self) -> range:
         return range(self.size)
 
     @property
-    def data(self) -> range:
-        return range(self.size)
+    def data(self) -> np.ndarray:
+        return np.arange(self.size, dtype=int)
 
     def __array__(self) -> np.ndarray:
         return np.arange(self.size)
@@ -517,6 +517,7 @@ class GroupBy(Generic[T_Xarray]):
         "_dims",
         "_sizes",
         "_len",
+        "_by_chunked",
         # Save unstacked object for flox
         "_original_obj",
         "_codes",

xarray/groupers.py

@@ -13,6 +13,7 @@
 
 import numpy as np
 import pandas as pd
+from numpy.typing import ArrayLike
 
 from xarray.coding.cftime_offsets import BaseCFTimeOffset, _new_to_legacy_freq
 from xarray.core import duck_array_ops
@@ -99,12 +100,18 @@ def __init__(
         assert isinstance(full_index, pd.Index)
         self.full_index = full_index
 
-        if group_indices is None and not is_chunked_array(codes.data):
-            self.group_indices = tuple(
-                g
-                for g in _codes_to_group_indices(codes.data.ravel(), len(full_index))
-                if g
-            )
+        if group_indices is None:
+            if not is_chunked_array(codes.data):
+                self.group_indices = tuple(
+                    g
+                    for g in _codes_to_group_indices(
+                        codes.data.ravel(), len(full_index)
+                    )
+                    if g
+                )
+            else:
+                # We will not use this when grouping by a chunked array
+                self.group_indices = tuple()
         else:
             self.group_indices = group_indices
 
@@ -168,13 +175,11 @@ class UniqueGrouper(Grouper):
     """
 
     _group_as_index: pd.Index | None = field(default=None, repr=False)
-    labels: np.ndarray | None = field(default=None)
+    labels: ArrayLike | None = field(default=None)
 
     @property
     def group_as_index(self) -> pd.Index:
         """Caches the group DataArray as a pandas Index."""
-        if is_chunked_array(self.group):
-            raise ValueError("Please call compute manually.")
         if self._group_as_index is None:
             if self.group.ndim == 1:
                 self._group_as_index = self.group.to_index()
@@ -214,7 +219,7 @@ def _factorize_given_labels(self, group: T_Group) -> EncodedGroups:
         )
         return EncodedGroups(
             codes=codes,
-            full_index=pd.Index(self.labels),
+            full_index=pd.Index(self.labels),  # type: ignore[arg-type]
             unique_coord=Variable(
                 dims=codes.name,
                 data=self.labels,
@@ -332,7 +337,7 @@ def __post_init__(self) -> None:
             raise ValueError("All bin edges are NaN.")
 
     def _cut(self, data):
-        return pd.cut(  # type: ignore [call-overload]
+        return pd.cut(
             np.asarray(data).ravel(),
             bins=self.bins,
             right=self.right,