Skip to content

Commit

Permalink
Fix typing
Browse files Browse the repository at this point in the history
  • Loading branch information
dcherian committed Mar 10, 2023
1 parent f46f5d9 commit 94aca25
Showing 1 changed file with 11 additions and 17 deletions.
28 changes: 11 additions & 17 deletions xarray/core/groupby.py
Original file line number Diff line number Diff line change
Expand Up @@ -288,7 +288,7 @@ def _apply_loffset(


class Grouper:
def __init__(self, group):
def __init__(self, group: T_Group):
self.group = group
self.codes = None
self.labels = None
Expand All @@ -298,21 +298,21 @@ def __init__(self, group):
self._group_as_index = None

@property
def name(self):
def name(self) -> Hashable:
return self.group1d.name

@property
def size(self):
def size(self) -> int:
return len(self)

def __len__(self):
def __len__(self) -> int:
return len(self.full_index)

@property
def dims(self):
return self.group1d.dims

def factorize(self, squeeze):
def factorize(self, squeeze: bool) -> None:
raise NotImplementedError

@property
Expand All @@ -328,7 +328,7 @@ def group_as_index(self) -> pd.Index:
self._group_as_index = safe_cast_to_index(self.group1d)
return self._group_as_index

def _resolve_group(self, obj):
def _resolve_group(self, obj) -> None:
from xarray.core.dataarray import DataArray

group = self.group
Expand Down Expand Up @@ -369,14 +369,14 @@ def _resolve_group(self, obj):


class UniqueGrouper(Grouper):
def factorize(self, squeeze):
def factorize(self, squeeze) -> None:
is_dimension = self.group.dims == (self.group.name,)
if is_dimension and self.is_unique_and_monotonic:
self._factorize_dummy(squeeze)
else:
self._factorize_unique()

def _factorize_unique(self):
def _factorize_unique(self) -> None:
# look through group to find the unique values
sort = not isinstance(self.group_as_index, pd.MultiIndex)
unique_values, group_indices, codes = unique_value_groups(
Expand All @@ -393,7 +393,7 @@ def _factorize_unique(self):
self.group_indices = group_indices
self.full_index = self.unique_coord

def _factorize_dummy(self, squeeze):
def _factorize_dummy(self, squeeze) -> None:
size = self.group.size
# no need to factorize
if not squeeze:
Expand Down Expand Up @@ -423,7 +423,7 @@ def __init__(self, group, bins, cut_kwargs):
self.bins = bins
self.cut_kwargs = cut_kwargs

def factorize(self, squeeze):
def factorize(self, squeeze: bool) -> None:
from xarray.core.dataarray import DataArray

data = self.group1d.values
Expand Down Expand Up @@ -526,13 +526,7 @@ def first_items(self):
_apply_loffset(self.loffset, first_items)
return first_items, codes

def factorize(
self, squeeze
) -> tuple[
DataArray | IndexVariable | _DummyGroup,
list[slice] | list[list[int]] | np.ndarray,
np.ndarray,
]:
def factorize(self, squeeze: bool) -> None:
self.full_index, first_items, codes = self._get_index_and_items()
sbins = first_items.values.astype(np.int64)
self.group_indices = [slice(i, j) for i, j in zip(sbins[:-1], sbins[1:])] + [
Expand Down

0 comments on commit 94aca25

Please sign in to comment.