Skip to content

Commit

Permalink
Typing fixes
Browse files Browse the repository at this point in the history
  • Loading branch information
dcherian committed Sep 20, 2024
1 parent 93e786b commit dfdc96a
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 13 deletions.
5 changes: 3 additions & 2 deletions xarray/core/groupby.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -517,6 +517,7 @@ class GroupBy(Generic[T_Xarray]):
"_dims",
"_sizes",
"_len",
"_by_chunked",
# Save unstacked object for flox
"_original_obj",
"_codes",
Expand Down
27 changes: 16 additions & 11 deletions xarray/groupers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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()
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down

0 comments on commit dfdc96a

Please sign in to comment.