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))