|
9 | 9 | from dataclasses import dataclass, field
|
10 | 10 | from datetime import timedelta
|
11 | 11 | from html import escape
|
12 |
| -from typing import TYPE_CHECKING, Any, Callable |
| 12 | +from typing import TYPE_CHECKING, Any, Callable, overload |
13 | 13 |
|
14 | 14 | import numpy as np
|
15 | 15 | import pandas as pd
|
@@ -236,14 +236,34 @@ def expanded_indexer(key, ndim):
|
236 | 236 | return tuple(new_key)
|
237 | 237 |
|
238 | 238 |
|
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)) |
241 | 252 |
|
242 | 253 |
|
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) |
247 | 267 |
|
248 | 268 |
|
249 | 269 | def slice_slice(old_slice: slice, applied_slice: slice, size: int) -> slice:
|
@@ -316,11 +336,15 @@ def __repr__(self) -> str:
|
316 | 336 | return f"{type(self).__name__}({self.tuple})"
|
317 | 337 |
|
318 | 338 |
|
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: |
320 | 344 | return None if value is None else operator.index(value)
|
321 | 345 |
|
322 | 346 |
|
323 |
| -def as_integer_slice(value): |
| 347 | +def as_integer_slice(value: slice) -> slice: |
324 | 348 | start = as_integer_or_none(value.start)
|
325 | 349 | stop = as_integer_or_none(value.stop)
|
326 | 350 | step = as_integer_or_none(value.step)
|
|
0 commit comments