1import warnings
2
3from numpy import ndarray, asarray, integer, bool_, array2string, empty, intp
4
5from .ndindex import NDIndex, asshape
6
7class ArrayIndex(NDIndex):
8    """
9    Superclass for array indices
10
11    This class should not be instantiated directly. Rather, use one of its
12    subclasses, :class:`~.IntegerArray` or :class:`~.BooleanArray`.
13
14    To subclass this, define the `dtype` attribute, as well as all the usual
15    ndindex methods.
16    """
17    __slots__ = ()
18
19    # Subclasses should redefine this
20    dtype = None
21
22    def _typecheck(self, idx, shape=None, _copy=True):
23        if self.dtype is None:
24            raise TypeError("Do not instantiate the superclass ArrayIndex directly")
25
26        if shape is not None:
27            if idx != []:
28                raise ValueError("The shape argument is only allowed for empty arrays (idx=[])")
29            shape = asshape(shape)
30            if 0 not in shape:
31                raise ValueError("The shape argument must be an empty shape")
32            idx = empty(shape, dtype=self.dtype)
33
34        if isinstance(idx, (list, ndarray, bool, integer, int, bool_)):
35            # Ignore deprecation warnings for things like [1, []]. These will be
36            # filtered out anyway since they produce object arrays.
37            with warnings.catch_warnings(record=True):
38                a = asarray(idx)
39                if a is idx and _copy:
40                    a = a.copy()
41                if isinstance(idx, list) and 0 in a.shape:
42                    if not _copy:
43                        raise ValueError("_copy=False is not allowed with list input")
44                    a = a.astype(self.dtype)
45            if self.dtype == intp and issubclass(a.dtype.type, integer):
46                if a.dtype != self.dtype:
47                    if not _copy:
48                        raise ValueError("If _copy=False, the input array dtype must already be intp")
49                    a = a.astype(self.dtype)
50            if a.dtype != self.dtype:
51                raise TypeError(f"The input array to {self.__class__.__name__} must have dtype {self.dtype.__name__}, not {a.dtype}")
52            a.flags.writeable = False
53            return (a,)
54        raise TypeError(f"{self.__class__.__name__} must be created with an array with dtype {self.dtype.__name__}")
55
56    # These will allow array == ArrayIndex to give True or False instead of
57    # returning an array.
58    __array_ufunc__ = None
59    def __array_function__(self, func, types, args, kwargs):
60        return NotImplemented
61
62    def __array__(self):
63        raise TypeError(f"Cannot convert {self.__class__.__name__} to an array. Use .array instead.")
64
65
66    @property
67    def raw(self):
68        return self.args[0]
69
70    @property
71    def array(self):
72        """
73        Return the NumPy array of self.
74
75        This is the same as `self.args[0]`.
76
77        >>> from ndindex import IntegerArray, BooleanArray
78        >>> IntegerArray([0, 1]).array
79        array([0, 1])
80        >>> BooleanArray([False, True]).array
81        array([False, True])
82
83        """
84        return self.args[0]
85
86    @property
87    def shape(self):
88        """
89        Return the shape of the array of self.
90
91        This is the same as `self.array.shape`. Note that this is **not** the
92        same as the shape of an array that is indexed by `self`. Use
93        :meth:`~.NDIndex.newshape` to get that.
94
95        >>> from ndindex import IntegerArray, BooleanArray
96        >>> IntegerArray([[0], [1]]).shape
97        (2, 1)
98        >>> BooleanArray([[False], [True]]).shape
99        (2, 1)
100
101        """
102        return self.array.shape
103
104    @property
105    def ndim(self):
106        """
107        Return the number of dimensions of the array of self.
108
109        This is the same as `self.array.ndim`. Note that this is **not** the
110        same as the number of dimensions of an array that is indexed by
111        `self`. Use `len` on :meth:`~.NDIndex.newshape` to get that.
112
113        >>> from ndindex import IntegerArray, BooleanArray
114        >>> IntegerArray([[0], [1]]).ndim
115        2
116        >>> BooleanArray([[False], [True]]).ndim
117        2
118
119        """
120        return self.array.ndim
121
122    @property
123    def size(self):
124        """
125        Return the number of elements of the array of self.
126
127        This is the same as `self.array.size`. Note that this is **not** the
128        same as the number of elements of an array that is indexed by `self`.
129        Use `np.prod` on :meth:`~.NDIndex.newshape` to get that.
130
131        >>> from ndindex import IntegerArray, BooleanArray
132        >>> IntegerArray([[0], [1]]).size
133        2
134        >>> BooleanArray([[False], [True]]).size
135        2
136
137        """
138        return self.array.size
139
140    # The repr form recreates the object. The str form gives the truncated
141    # array string and is explicitly non-valid Python (doesn't have commas).
142    def __repr__(self):
143        if 0 not in self.shape:
144            arg = repr(self.array.tolist())
145        else:
146            arg = f"[], shape={self.shape}"
147        return f"{self.__class__.__name__}({arg})"
148
149    def __str__(self):
150        return (self.__class__.__name__
151                + "("
152                + array2string(self.array).replace('\n', '')
153                + ")")
154
155    def __hash__(self):
156        return hash(self.array.tobytes())
157