1"""
2Low-dependency indexing utilities.
3"""
4
5from __future__ import annotations
6
7from typing import (
8 TYPE_CHECKING,
9 Any,
10)
11
12import numpy as np
13
14from pandas._libs import lib
15from pandas.util._decorators import set_module
16
17from pandas.core.dtypes.common import (
18 is_array_like_deprecate_non_pandas,
19 is_bool_dtype,
20 is_integer,
21 is_integer_dtype,
22 is_list_like,
23)
24from pandas.core.dtypes.dtypes import ExtensionDtype
25from pandas.core.dtypes.generic import (
26 ABCIndex,
27 ABCSeries,
28)
29
30if TYPE_CHECKING:
31 from pandas._typing import AnyArrayLike
32
33 from pandas.core.frame import DataFrame
34 from pandas.core.indexes.base import Index
35
36# -----------------------------------------------------------
37# Indexer Identification
38
39
40def is_valid_positional_slice(slc: slice) -> bool:
41 """
42 Check if a slice object can be interpreted as a positional indexer.
43
44 Parameters
45 ----------
46 slc : slice
47
48 Returns
49 -------
50 bool
51
52 Notes
53 -----
54 A valid positional slice may also be interpreted as a label-based slice
55 depending on the index being sliced.
56 """
57 return (
58 lib.is_int_or_none(slc.start)
59 and lib.is_int_or_none(slc.stop)
60 and lib.is_int_or_none(slc.step)
61 )
62
63
64def is_list_like_indexer(key) -> bool:
65 """
66 Check if we have a list-like indexer that is *not* a NamedTuple.
67
68 Parameters
69 ----------
70 key : object
71
72 Returns
73 -------
74 bool
75 """
76 # allow a list_like, but exclude NamedTuples which can be indexers
77 return is_list_like(key) and not (isinstance(key, tuple) and type(key) is not tuple)
78
79
80def is_scalar_indexer(indexer, ndim: int) -> bool:
81 """
82 Return True if we are all scalar indexers.
83
84 Parameters
85 ----------
86 indexer : object
87 ndim : int
88 Number of dimensions in the object being indexed.
89
90 Returns
91 -------
92 bool
93 """
94 if ndim == 1 and is_integer(indexer):
95 # GH37748: allow indexer to be an integer for Series
96 return True
97 if isinstance(indexer, tuple) and len(indexer) == ndim:
98 return all(is_integer(x) for x in indexer)
99 return False
100
101
102def is_empty_indexer(indexer) -> bool:
103 """
104 Check if we have an empty indexer.
105
106 Parameters
107 ----------
108 indexer : object
109
110 Returns
111 -------
112 bool
113 """
114 if is_list_like(indexer) and not len(indexer):
115 return True
116 if not isinstance(indexer, tuple):
117 indexer = (indexer,)
118 return any(isinstance(idx, np.ndarray) and len(idx) == 0 for idx in indexer)
119
120
121# -----------------------------------------------------------
122# Indexer Validation
123
124
125def check_setitem_lengths(indexer, value, values) -> bool:
126 """
127 Validate that value and indexer are the same length.
128
129 A special-case is allowed for when the indexer is a boolean array
130 and the number of true values equals the length of ``value``. In
131 this case, no exception is raised.
132
133 Parameters
134 ----------
135 indexer : sequence
136 Key for the setitem.
137 value : array-like
138 Value for the setitem.
139 values : array-like
140 Values being set into.
141
142 Returns
143 -------
144 bool
145 Whether this is an empty listlike setting which is a no-op.
146
147 Raises
148 ------
149 ValueError
150 When the indexer is an ndarray or list and the lengths don't match.
151 """
152 no_op = False
153
154 if isinstance(indexer, (np.ndarray, list)):
155 # We can ignore other listlikes because they are either
156 # a) not necessarily 1-D indexers, e.g. tuple
157 # b) boolean indexers e.g. BoolArray
158 if is_list_like(value):
159 if len(indexer) != len(value) and values.ndim == 1:
160 # boolean with truth values == len of the value is ok too
161 if isinstance(indexer, list):
162 indexer = np.array(indexer)
163 if not (
164 isinstance(indexer, np.ndarray)
165 and indexer.dtype == np.bool_
166 and indexer.sum() == len(value)
167 ):
168 raise ValueError(
169 "cannot set using a list-like indexer "
170 "with a different length than the value"
171 )
172 if not len(indexer):
173 no_op = True
174
175 elif isinstance(indexer, slice):
176 if is_list_like(value):
177 if len(value) != length_of_indexer(indexer, values) and values.ndim == 1:
178 # In case of two dimensional value is used row-wise and broadcasted
179 raise ValueError(
180 "cannot set using a slice indexer with a "
181 "different length than the value"
182 )
183 if not len(value):
184 no_op = True
185
186 return no_op
187
188
189def validate_indices(indices: np.ndarray, n: int) -> None:
190 """
191 Perform bounds-checking for an indexer.
192
193 -1 is allowed for indicating missing values.
194
195 Parameters
196 ----------
197 indices : ndarray
198 n : int
199 Length of the array being indexed.
200
201 Raises
202 ------
203 ValueError
204
205 Examples
206 --------
207 >>> validate_indices(np.array([1, 2]), 3) # OK
208
209 >>> validate_indices(np.array([1, -2]), 3)
210 Traceback (most recent call last):
211 ...
212 ValueError: negative dimensions are not allowed
213
214 >>> validate_indices(np.array([1, 2, 3]), 3)
215 Traceback (most recent call last):
216 ...
217 IndexError: indices are out-of-bounds
218
219 >>> validate_indices(np.array([-1, -1]), 0) # OK
220
221 >>> validate_indices(np.array([0, 1]), 0)
222 Traceback (most recent call last):
223 ...
224 IndexError: indices are out-of-bounds
225 """
226 if len(indices):
227 min_idx = indices.min()
228 if min_idx < -1:
229 msg = f"'indices' contains values less than allowed ({min_idx} < -1)"
230 raise ValueError(msg)
231
232 max_idx = indices.max()
233 if max_idx >= n:
234 raise IndexError("indices are out-of-bounds")
235
236
237# -----------------------------------------------------------
238# Indexer Conversion
239
240
241def maybe_convert_indices(indices, n: int, verify: bool = True) -> np.ndarray:
242 """
243 Attempt to convert indices into valid, positive indices.
244
245 If we have negative indices, translate to positive here.
246 If we have indices that are out-of-bounds, raise an IndexError.
247
248 Parameters
249 ----------
250 indices : array-like
251 Array of indices that we are to convert.
252 n : int
253 Number of elements in the array that we are indexing.
254 verify : bool, default True
255 Check that all entries are between 0 and n - 1, inclusive.
256
257 Returns
258 -------
259 array-like
260 An array-like of positive indices that correspond to the ones
261 that were passed in initially to this function.
262
263 Raises
264 ------
265 IndexError
266 One of the converted indices either exceeded the number of,
267 elements (specified by `n`), or was still negative.
268 """
269 if isinstance(indices, list):
270 indices = np.array(indices)
271 if len(indices) == 0:
272 # If `indices` is empty, np.array will return a float,
273 # and will cause indexing errors.
274 return np.empty(0, dtype=np.intp)
275
276 mask = indices < 0
277 if mask.any():
278 indices = indices.copy()
279 indices[mask] += n
280
281 if verify:
282 mask = (indices >= n) | (indices < 0)
283 if mask.any():
284 raise IndexError("indices are out-of-bounds")
285 return indices
286
287
288# -----------------------------------------------------------
289# Unsorted
290
291
292def length_of_indexer(indexer, target=None) -> int:
293 """
294 Return the expected length of target[indexer]
295
296 Returns
297 -------
298 int
299 """
300 if target is not None and isinstance(indexer, slice):
301 target_len = len(target)
302 # GH#66100 slice.indices() correctly resolves None bounds for both
303 # positive and negative step, which the previous manual arithmetic
304 # did not (it could return negative lengths for negative-step
305 # slices with start/stop left as None).
306 start, stop, step = indexer.indices(target_len)
307 return len(range(start, stop, step))
308 elif isinstance(indexer, (ABCSeries, ABCIndex, np.ndarray, list)):
309 if isinstance(indexer, list):
310 indexer = np.array(indexer)
311
312 if indexer.dtype == bool:
313 # GH#25774
314 # np.asarray because Index has no .sum; GH#68021
315 return int(np.asarray(indexer).sum())
316 return len(indexer)
317 elif isinstance(indexer, range):
318 try:
319 return len(indexer)
320 except OverflowError:
321 step = indexer.step
322 if step > 0:
323 low, high = indexer.start, indexer.stop
324 else:
325 low, high = indexer.stop, indexer.start
326 step = -step
327 if low >= high:
328 return 0
329 return (high - low - 1) // step + 1
330 elif not is_list_like_indexer(indexer):
331 return 1
332 raise AssertionError("cannot find the length of the indexer")
333
334
335def disallow_ndim_indexing(result) -> None:
336 """
337 Helper function to disallow multi-dimensional indexing on 1D Series/Index.
338
339 GH#27125 indexer like idx[:, None] expands dim, but we cannot do that
340 and keep an index, so we used to return ndarray, which was deprecated
341 in GH#30588.
342 """
343 if np.ndim(result) > 1:
344 raise ValueError(
345 "Multi-dimensional indexing (e.g. `obj[:, None]`) is no longer "
346 "supported. Convert to a numpy array before indexing instead."
347 )
348
349
350def unpack_1tuple(tup):
351 """
352 If we have a length-1 tuple/list that contains a slice, unpack to just
353 the slice.
354
355 Notes
356 -----
357 The list case is deprecated.
358 """
359 if len(tup) == 1 and isinstance(tup[0], slice):
360 # if we don't have a MultiIndex, we may still be able to handle
361 # a 1-tuple. see test_1tuple_without_multiindex
362
363 if isinstance(tup, list):
364 # GH#31299
365 raise ValueError(
366 "Indexing with a single-item list containing a "
367 "slice is not allowed. Pass a tuple instead.",
368 )
369
370 return tup[0]
371 return tup
372
373
374def check_key_length(columns: Index, key, value: DataFrame) -> None:
375 """
376 Checks if a key used as indexer has the same length as the columns it is
377 associated with.
378
379 Parameters
380 ----------
381 columns : Index The columns of the DataFrame to index.
382 key : A list-like of keys to index with.
383 value : DataFrame The value to set for the keys.
384
385 Raises
386 ------
387 ValueError: If the length of key is not equal to the number of columns in value
388 or if the number of columns referenced by key is not equal to number
389 of columns.
390 """
391 if columns.is_unique:
392 if len(value.columns) != len(key):
393 raise ValueError("Columns must be same length as key")
394 # Missing keys in columns are represented as -1
395 elif len(columns.get_indexer_non_unique(key)[0]) != len(value.columns):
396 raise ValueError("Columns must be same length as key")
397
398
399def unpack_tuple_and_ellipses(item: tuple[Any, ...]):
400 """
401 Possibly unpack arr[..., n] to arr[n]
402 """
403 if len(item) > 1:
404 # Note: we are assuming this indexing is being done on a 1D arraylike
405 if item[0] is Ellipsis:
406 item = item[1:]
407 elif item[-1] is Ellipsis:
408 item = item[:-1]
409
410 if len(item) > 1:
411 raise IndexError("too many indices for array.")
412
413 item = item[0]
414 return item
415
416
417def getitem_returns_view(arr, key) -> bool:
418 """
419 Check if an ``arr.__getitem__`` call with given ``key`` would return a view
420 or not.
421 """
422 if not isinstance(key, tuple):
423 key = (key,)
424
425 # filter out Ellipsis and np.newaxis
426 key = tuple(k for k in key if k is not Ellipsis and k is not np.newaxis)
427 if not key:
428 return True
429 # single integer gives view if selecting subset of 2D array
430 if arr.ndim == 2 and lib.is_integer(key[0]):
431 return True
432 # slices always give views
433 if all(isinstance(k, slice) for k in key):
434 return True
435 return False
436
437
438# -----------------------------------------------------------
439# Public indexer validation
440
441
442@set_module("pandas.api.indexers")
443def check_array_indexer(array: AnyArrayLike, indexer: Any) -> Any:
444 """
445 Check if `indexer` is a valid array indexer for `array`.
446
447 For a boolean mask, `array` and `indexer` are checked to have the same
448 length. The dtype is validated, and if it is an integer or boolean
449 ExtensionArray, it is checked if there are missing values present, and
450 it is converted to the appropriate numpy array. Other dtypes will raise
451 an error.
452
453 Non-array indexers (integer, slice, Ellipsis, tuples, ..) are passed
454 through as is.
455
456 Parameters
457 ----------
458 array : array-like
459 The array that is being indexed (only used for the length).
460 indexer : array-like, list-like, int, slice, or other indexer
461 The indexer used for indexing. Array-like and list-like inputs that
462 are not yet a numpy array or an ExtensionArray are converted to one.
463 Non-array indexers (int, slice, Ellipsis, tuples, etc.) are passed
464 through as is.
465
466 Returns
467 -------
468 numpy.ndarray
469 The validated indexer as a numpy array that can be used to index.
470
471 Raises
472 ------
473 IndexError
474 When the lengths don't match.
475 ValueError
476 When `indexer` cannot be converted to a numpy ndarray to index
477 (e.g. presence of missing values).
478
479 See Also
480 --------
481 api.types.is_bool_dtype : Check if `key` is of boolean dtype.
482
483 Examples
484 --------
485 When checking a boolean mask, a boolean ndarray is returned when the
486 arguments are all valid.
487
488 >>> mask = pd.array([True, False])
489 >>> arr = pd.array([1, 2])
490 >>> pd.api.indexers.check_array_indexer(arr, mask)
491 array([ True, False])
492
493 An IndexError is raised when the lengths don't match.
494
495 >>> mask = pd.array([True, False, True])
496 >>> pd.api.indexers.check_array_indexer(arr, mask)
497 Traceback (most recent call last):
498 ...
499 IndexError: Boolean index has wrong length: 3 instead of 2.
500
501 NA values in a boolean array are treated as False.
502
503 >>> mask = pd.array([True, pd.NA])
504 >>> pd.api.indexers.check_array_indexer(arr, mask)
505 array([ True, False])
506
507 A numpy boolean mask will get passed through (if the length is correct):
508
509 >>> mask = np.array([True, False])
510 >>> pd.api.indexers.check_array_indexer(arr, mask)
511 array([ True, False])
512
513 Integer and slice indexers are passed through as is:
514
515 >>> pd.api.indexers.check_array_indexer(arr, 1)
516 1
517 >>> pd.api.indexers.check_array_indexer(arr, slice(0, 1, 1))
518 slice(0, 1, 1)
519
520 Similarly for integer indexers, an integer ndarray is returned when it is
521 a valid indexer, otherwise an error is (for integer indexers, a matching
522 length is not required):
523
524 >>> indexer = pd.array([0, 2], dtype="Int64")
525 >>> arr = pd.array([1, 2, 3])
526 >>> pd.api.indexers.check_array_indexer(arr, indexer)
527 array([0, 2])
528
529 >>> indexer = pd.array([0, pd.NA], dtype="Int64")
530 >>> pd.api.indexers.check_array_indexer(arr, indexer)
531 Traceback (most recent call last):
532 ...
533 ValueError: Cannot index with an integer indexer containing NA values
534
535 For non-integer/boolean dtypes, an appropriate error is raised:
536
537 >>> indexer = np.array([0.0, 2.0], dtype="float64")
538 >>> pd.api.indexers.check_array_indexer(arr, indexer)
539 Traceback (most recent call last):
540 ...
541 IndexError: arrays used as indices must be of integer or boolean type
542 """
543 from pandas.core.construction import array as pd_array
544
545 # whatever is not an array-like is returned as-is (possible valid array
546 # indexers that are not array-like: integer, slice, Ellipsis, None)
547 # In this context, tuples are not considered as array-like, as they have
548 # a specific meaning in indexing (multi-dimensional indexing)
549 if is_list_like(indexer):
550 if isinstance(indexer, tuple):
551 return indexer
552 else:
553 return indexer
554
555 # convert list-likes to array
556 if not is_array_like_deprecate_non_pandas(indexer):
557 indexer = pd_array(indexer)
558 if len(indexer) == 0:
559 # empty list is converted to float array by pd.array
560 indexer = np.array([], dtype=np.intp)
561
562 dtype = indexer.dtype
563 if is_bool_dtype(dtype):
564 if isinstance(dtype, ExtensionDtype):
565 indexer = indexer.to_numpy(dtype=bool, na_value=False)
566 else:
567 indexer = np.asarray(indexer, dtype=bool)
568
569 # GH26658
570 if len(indexer) != len(array):
571 raise IndexError(
572 f"Boolean index has wrong length: "
573 f"{len(indexer)} instead of {len(array)}"
574 )
575 elif is_integer_dtype(dtype):
576 try:
577 indexer = np.asarray(indexer, dtype=np.intp)
578 except ValueError as err:
579 raise ValueError(
580 "Cannot index with an integer indexer containing NA values"
581 ) from err
582 else:
583 raise IndexError("arrays used as indices must be of integer or boolean type")
584
585 return indexer