Skip to content

Commit

Permalink
Some typing (#190)
Browse files Browse the repository at this point in the history
* Some typing

* fixes.
  • Loading branch information
dcherian authored Nov 15, 2022
1 parent 9b01c48 commit 51fb6e9
Showing 1 changed file with 5 additions and 5 deletions.
10 changes: 5 additions & 5 deletions flox/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -420,7 +420,7 @@ def offset_labels(labels: np.ndarray, ngroups: int) -> tuple[np.ndarray, int]:
def factorize_(
by: tuple,
axis: T_AxesOpt,
expected_groups: tuple[pd.Index, ...] = None,
expected_groups: tuple[pd.Index, ...] | None = None,
reindex: bool = False,
sort=True,
fastpath=False,
Expand Down Expand Up @@ -873,7 +873,7 @@ def _simple_combine(
return results


def _conc2(x_chunk, key1, key2=slice(None), axis: T_Axes = None) -> np.ndarray:
def _conc2(x_chunk, key1, key2=slice(None), axis: T_Axes | None = None) -> np.ndarray:
"""copied from dask.array.reductions.mean_combine"""
from dask.array.core import _concatenate2
from dask.utils import deepmap
Expand Down Expand Up @@ -1071,7 +1071,7 @@ def _reduce_blockwise(
return result


def _normalize_indexes(array, flatblocks, blkshape):
def _normalize_indexes(array: DaskArray, flatblocks, blkshape) -> tuple:
"""
.blocks accessor can only accept one iterable at a time,
but can handle multiple slices.
Expand All @@ -1083,7 +1083,7 @@ def _normalize_indexes(array, flatblocks, blkshape):
"""
unraveled = np.unravel_index(flatblocks, blkshape)

normalized: list[Union[int, np.ndarray, slice]] = []
normalized: list[Union[int, slice, list[int]]] = []
for ax, idx in enumerate(unraveled):
i = _unique(idx).squeeze()
if i.ndim == 0:
Expand Down Expand Up @@ -1397,7 +1397,7 @@ def dask_groupby_agg(
return (result, groups)


def _collapse_blocks_along_axes(reduced, axis, group_chunks):
def _collapse_blocks_along_axes(reduced: DaskArray, axis: T_Axes, group_chunks) -> DaskArray:
import dask.array
from dask.highlevelgraph import HighLevelGraph

Expand Down

0 comments on commit 51fb6e9

Please sign in to comment.