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