Skip to content

Commit a07e16c

Browse files
authored
Add typing to some functions in indexing.py (#8922)
* Add typing * Update indexing.py
1 parent 07c7f96 commit a07e16c

File tree

1 file changed

+33
-9
lines changed

1 file changed

+33
-9
lines changed

xarray/core/indexing.py

Lines changed: 33 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
from dataclasses import dataclass, field
1010
from datetime import timedelta
1111
from html import escape
12-
from typing import TYPE_CHECKING, Any, Callable
12+
from typing import TYPE_CHECKING, Any, Callable, overload
1313

1414
import numpy as np
1515
import pandas as pd
@@ -236,14 +236,34 @@ def expanded_indexer(key, ndim):
236236
return tuple(new_key)
237237

238238

239-
def _expand_slice(slice_, size: int) -> np.ndarray:
240-
return np.arange(*slice_.indices(size))
239+
def _normalize_slice(sl: slice, size: int) -> slice:
240+
"""
241+
Ensure that given slice only contains positive start and stop values
242+
(stop can be -1 for full-size slices with negative steps, e.g. [-10::-1])
243+
244+
Examples
245+
--------
246+
>>> _normalize_slice(slice(0, 9), 10)
247+
slice(0, 9, 1)
248+
>>> _normalize_slice(slice(0, -1), 10)
249+
slice(0, 9, 1)
250+
"""
251+
return slice(*sl.indices(size))
241252

242253

243-
def _normalize_slice(sl: slice, size) -> slice:
244-
"""Ensure that given slice only contains positive start and stop values
245-
(stop can be -1 for full-size slices with negative steps, e.g. [-10::-1])"""
246-
return slice(*sl.indices(size))
254+
def _expand_slice(slice_: slice, size: int) -> np.ndarray[Any, np.dtype[np.integer]]:
255+
"""
256+
Expand slice to an array containing only positive integers.
257+
258+
Examples
259+
--------
260+
>>> _expand_slice(slice(0, 9), 10)
261+
array([0, 1, 2, 3, 4, 5, 6, 7, 8])
262+
>>> _expand_slice(slice(0, -1), 10)
263+
array([0, 1, 2, 3, 4, 5, 6, 7, 8])
264+
"""
265+
sl = _normalize_slice(slice_, size)
266+
return np.arange(sl.start, sl.stop, sl.step)
247267

248268

249269
def slice_slice(old_slice: slice, applied_slice: slice, size: int) -> slice:
@@ -316,11 +336,15 @@ def __repr__(self) -> str:
316336
return f"{type(self).__name__}({self.tuple})"
317337

318338

319-
def as_integer_or_none(value):
339+
@overload
340+
def as_integer_or_none(value: int) -> int: ...
341+
@overload
342+
def as_integer_or_none(value: None) -> None: ...
343+
def as_integer_or_none(value: int | None) -> int | None:
320344
return None if value is None else operator.index(value)
321345

322346

323-
def as_integer_slice(value):
347+
def as_integer_slice(value: slice) -> slice:
324348
start = as_integer_or_none(value.start)
325349
stop = as_integer_or_none(value.stop)
326350
step = as_integer_or_none(value.step)

0 commit comments

Comments
 (0)