|
1 | 1 | from __future__ import annotations
|
2 | 2 |
|
3 | 3 | import itertools
|
4 |
| -from typing import Any |
| 4 | +from typing import Any, Union |
5 | 5 |
|
6 | 6 | import numpy as np
|
7 | 7 | import pandas as pd
|
@@ -642,7 +642,7 @@ def setup(self):
|
642 | 642 |
|
643 | 643 | def test_arrayize_vectorized_indexer(self) -> None:
|
644 | 644 | for i, j, k in itertools.product(self.indexers, repeat=3):
|
645 |
| - vindex = indexing.VectorizedIndexer((i, j, k)) |
| 645 | + vindex = indexing.VectorizedIndexer((i, j, k)) # type: ignore[arg-type] |
646 | 646 | vindex_array = indexing._arrayize_vectorized_indexer(
|
647 | 647 | vindex, self.data.shape
|
648 | 648 | )
|
@@ -676,46 +676,58 @@ def test_arrayize_vectorized_indexer(self) -> None:
|
676 | 676 | np.testing.assert_array_equal(b, np.arange(5)[:, np.newaxis])
|
677 | 677 |
|
678 | 678 |
|
679 |
| -def get_indexers(shape, mode): |
| 679 | +def get_indexers( |
| 680 | + shape: tuple[int, ...], mode: str |
| 681 | +) -> Union[indexing.VectorizedIndexer, indexing.OuterIndexer, indexing.BasicIndexer]: |
680 | 682 | if mode == "vectorized":
|
681 | 683 | indexed_shape = (3, 4)
|
682 |
| - indexer = tuple(np.random.randint(0, s, size=indexed_shape) for s in shape) |
683 |
| - return indexing.VectorizedIndexer(indexer) |
| 684 | + indexer_v = tuple(np.random.randint(0, s, size=indexed_shape) for s in shape) |
| 685 | + return indexing.VectorizedIndexer(indexer_v) |
684 | 686 |
|
685 | 687 | elif mode == "outer":
|
686 |
| - indexer = tuple(np.random.randint(0, s, s + 2) for s in shape) |
687 |
| - return indexing.OuterIndexer(indexer) |
| 688 | + indexer_o = tuple(np.random.randint(0, s, s + 2) for s in shape) |
| 689 | + return indexing.OuterIndexer(indexer_o) |
688 | 690 |
|
689 | 691 | elif mode == "outer_scalar":
|
690 |
| - indexer = (np.random.randint(0, 3, 4), 0, slice(None, None, 2)) |
691 |
| - return indexing.OuterIndexer(indexer[: len(shape)]) |
| 692 | + indexer_os: tuple[Any, ...] = ( |
| 693 | + np.random.randint(0, 3, 4), |
| 694 | + 0, |
| 695 | + slice(None, None, 2), |
| 696 | + ) |
| 697 | + return indexing.OuterIndexer(indexer_os[: len(shape)]) |
692 | 698 |
|
693 | 699 | elif mode == "outer_scalar2":
|
694 |
| - indexer = (np.random.randint(0, 3, 4), -2, slice(None, None, 2)) |
695 |
| - return indexing.OuterIndexer(indexer[: len(shape)]) |
| 700 | + indexer_os2: tuple[Any, ...] = ( |
| 701 | + np.random.randint(0, 3, 4), |
| 702 | + -2, |
| 703 | + slice(None, None, 2), |
| 704 | + ) |
| 705 | + return indexing.OuterIndexer(indexer_os2[: len(shape)]) |
696 | 706 |
|
697 | 707 | elif mode == "outer1vec":
|
698 |
| - indexer = [slice(2, -3) for s in shape] |
699 |
| - indexer[1] = np.random.randint(0, shape[1], shape[1] + 2) |
700 |
| - return indexing.OuterIndexer(tuple(indexer)) |
| 708 | + indexer_o1v: list[Any] = [slice(2, -3) for s in shape] |
| 709 | + indexer_o1v[1] = np.random.randint(0, shape[1], shape[1] + 2) |
| 710 | + return indexing.OuterIndexer(tuple(indexer_o1v)) |
701 | 711 |
|
702 | 712 | elif mode == "basic": # basic indexer
|
703 |
| - indexer = [slice(2, -3) for s in shape] |
704 |
| - indexer[0] = 3 |
705 |
| - return indexing.BasicIndexer(tuple(indexer)) |
| 713 | + indexer_b: list[Any] = [slice(2, -3) for s in shape] |
| 714 | + indexer_b[0] = 3 |
| 715 | + return indexing.BasicIndexer(tuple(indexer_b)) |
706 | 716 |
|
707 | 717 | elif mode == "basic1": # basic indexer
|
708 | 718 | return indexing.BasicIndexer((3,))
|
709 | 719 |
|
710 | 720 | elif mode == "basic2": # basic indexer
|
711 |
| - indexer = [0, 2, 4] |
712 |
| - return indexing.BasicIndexer(tuple(indexer[: len(shape)])) |
| 721 | + indexer_b2 = [0, 2, 4] |
| 722 | + return indexing.BasicIndexer(tuple(indexer_b2[: len(shape)])) |
713 | 723 |
|
714 | 724 | elif mode == "basic3": # basic indexer
|
715 |
| - indexer = [slice(None) for s in shape] |
716 |
| - indexer[0] = slice(-2, 2, -2) |
717 |
| - indexer[1] = slice(1, -1, 2) |
718 |
| - return indexing.BasicIndexer(tuple(indexer[: len(shape)])) |
| 725 | + indexer_b3: list[Any] = [slice(None) for s in shape] |
| 726 | + indexer_b3[0] = slice(-2, 2, -2) |
| 727 | + indexer_b3[1] = slice(1, -1, 2) |
| 728 | + return indexing.BasicIndexer(tuple(indexer_b3[: len(shape)])) |
| 729 | + |
| 730 | + raise ValueError(f"Unknown mode: {mode}") |
719 | 731 |
|
720 | 732 |
|
721 | 733 | @pytest.mark.parametrize("size", [100, 99])
|
|
0 commit comments