75 lines
2.5 KiB
Python
75 lines
2.5 KiB
Python
from typing import List, Tuple, Union # upip.install("pycopy-typing")
|
|
from ulab import numpy as np
|
|
|
|
_DType = int
|
|
_RClassKeyType = Union[slice, int, float, list, tuple, np.ndarray]
|
|
|
|
# this is a stripped down version of RClass (used by np.r_[...etc])
|
|
# it doesn't include support for string arguments as the first index element
|
|
class RClass:
|
|
|
|
def __getitem__(self, key: Union[_RClassKeyType, Tuple[_RClassKeyType, ...]]):
|
|
|
|
if not isinstance(key, tuple):
|
|
key = (key,)
|
|
|
|
objs: List[np.ndarray] = []
|
|
scalars: List[int] = []
|
|
arraytypes: List[_DType] = []
|
|
scalartypes: List[_DType] = []
|
|
|
|
# these may get overridden in following loop
|
|
axis = 0
|
|
|
|
for idx, item in enumerate(key):
|
|
scalar = False
|
|
|
|
try:
|
|
if isinstance(item, np.ndarray):
|
|
newobj = item
|
|
|
|
elif isinstance(item, slice):
|
|
step = item.step
|
|
start = item.start
|
|
stop = item.stop
|
|
if start is None:
|
|
start = 0
|
|
if step is None:
|
|
step = 1
|
|
if isinstance(step, complex):
|
|
size = int(abs(step))
|
|
newobj: np.ndarray = np.linspace(start, stop, num=size)
|
|
else:
|
|
newobj = np.arange(start, stop, step)
|
|
|
|
# if is number
|
|
elif isinstance(item, (int, float, bool)):
|
|
newobj = np.array([item])
|
|
scalars.append(len(objs))
|
|
scalar = True
|
|
scalartypes.append(newobj.dtype())
|
|
|
|
else:
|
|
newobj = np.array(item)
|
|
|
|
except TypeError:
|
|
raise Exception("index object %s of type %s is not supported by r_[]" % (
|
|
str(item), type(item)))
|
|
|
|
objs.append(newobj)
|
|
if not scalar and isinstance(newobj, np.ndarray):
|
|
arraytypes.append(newobj.dtype())
|
|
|
|
# Ensure that scalars won't up-cast unless warranted
|
|
final_dtype = min(arraytypes + scalartypes)
|
|
for idx, obj in enumerate(objs):
|
|
if obj.dtype != final_dtype:
|
|
objs[idx] = np.array(objs[idx], dtype=final_dtype)
|
|
|
|
return np.concatenate(tuple(objs), axis=axis)
|
|
|
|
# this seems weird - not sure what it's for
|
|
def __len__(self):
|
|
return 0
|
|
|
|
r_ = RClass()
|