def drop_duplicate_levels(index: tp.Index,
keep: tp.Optional[str] = None) -> tp.Index:
"""Drop levels in `index` with the same name and values.
Set `keep` to 'last' to keep last levels, otherwise 'first'.
Set `keep` to None to use the default."""
from vectorbt._settings import settings
broadcasting_cfg = settings['broadcasting']
if keep is None:
keep = broadcasting_cfg['keep']
if not isinstance(index, pd.MultiIndex):
return index
checks.assert_in(keep.lower(), ['first', 'last'])
levels = []
levels_to_drop = []
if keep == 'first':
r = range(0, index.nlevels)
else:
r = range(index.nlevels - 1, -1, -1) # loop backwards
for i in r:
level = (index.levels[i].name,
tuple(index.get_level_values(i).to_numpy().tolist()))
if level not in levels:
levels.append(level)
else:
levels_to_drop.append(i)
return index.droplevel(levels_to_drop)
def apply_on_index(self,
apply_func: tp.Callable,
*args,
axis: int = 1,
inplace: bool = False,
**kwargs) -> tp.Optional[tp.SeriesFrame]:
"""Apply function `apply_func` on index of the pandas object.
Set `axis` to 1 for columns and 0 for index.
If `inplace` is True, modifies the pandas object. Otherwise, returns a copy."""
checks.assert_in(axis, (0, 1))
if axis == 1:
obj_index = self.wrapper.columns
else:
obj_index = self.wrapper.index
obj_index = apply_func(obj_index, *args, **kwargs)
if inplace:
if axis == 1:
self.obj.columns = obj_index
else:
self.obj.index = obj_index
return None
else:
obj = self.obj.copy()
if axis == 1:
obj.columns = obj_index
else:
obj.index = obj_index
return obj
def apply_on_index(self, apply_func, *args, axis=1, inplace=False, **kwargs):
"""Apply function `apply_func` on index of the pandas object.
Set `axis` to 1 for columns and 0 for index.
If `inplace` is `True`, modifies the pandas object. Otherwise, returns a copy."""
checks.assert_in(axis, (0, 1))
if axis == 1:
obj_index = self.columns
else:
obj_index = self.index
obj_index = apply_func(obj_index, *args, **kwargs)
if inplace:
if axis == 1:
self._obj.columns = obj_index
else:
self._obj.index = obj_index
return None
else:
obj = self._obj.copy()
if axis == 1:
obj.columns = obj_index
else:
obj.index = obj_index
return obj
def _get_func_names(func_settings, setting, all_names):
func_input_names = func_settings.get(setting, None)
if func_input_names is None:
return []
else:
for name in func_input_names:
checks.assert_in(name, all_names)
return func_input_names
def _get_func_names(func_settings: tp.Kwargs, setting: str, all_names: tp.Sequence[str]) -> tp.List[str]:
func_input_names = func_settings.get(setting, None)
if func_input_names is None:
return []
else:
for name in func_input_names:
checks.assert_in(name, all_names)
return func_input_names
def get_index(arg, axis):
"""Get index of `arg` by `axis`."""
checks.assert_type(arg, (pd.Series, pd.DataFrame))
checks.assert_in(axis, (0, 1))
if axis == 0:
return arg.index
else:
if checks.is_series(arg):
if arg.name is not None:
return pd.Index([arg.name])
return pd.Index([0]) # same as how pandas does it
else:
return arg.columns
def __init__(self, wrapper, records_arr, idx_field=None):
Configured.__init__(self,
wrapper=wrapper,
records_arr=records_arr,
idx_field=idx_field)
checks.assert_type(wrapper, ArrayWrapper)
if not isinstance(records_arr, np.ndarray):
records_arr = np.asarray(records_arr)
checks.assert_not_none(records_arr.dtype.fields)
checks.assert_in('col', records_arr.dtype.names)
if idx_field is not None:
checks.assert_in(idx_field, records_arr.dtype.names)
else:
if 'idx' in records_arr.dtype.names:
idx_field = 'idx'
self._wrapper = wrapper
self._records_arr = records_arr
self._idx_field = idx_field
PandasIndexer.__init__(self, _records_indexing_func)
def __init__(self, wrapper, records_arr, idx_field='auto', **kwargs):
Wrapping.__init__(self,
wrapper,
records_arr=records_arr,
idx_field=idx_field,
**kwargs)
records_arr = np.asarray(records_arr)
checks.assert_not_none(records_arr.dtype.fields)
checks.assert_in('id', records_arr.dtype.names)
checks.assert_in('col', records_arr.dtype.names)
if idx_field == 'auto':
if 'idx' in records_arr.dtype.names:
idx_field = 'idx'
elif idx_field is not None:
checks.assert_in(idx_field, records_arr.dtype.names)
self._records_arr = records_arr
self._idx_field = idx_field
self._col_mapper = ColumnMapper(wrapper, records_arr['col'])
def test_assert_in(self):
checks.assert_in(0, (0, 1))
with pytest.raises(Exception) as e_info:
checks.assert_in(2, (0, 1))
