Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
def test_coldrop_one_col():
"""Testing the ColDrop pipeline stage."""
df = _test_df()
assert 'num1' in df.columns
stage = ColDrop('num1')
res_df = stage.apply(df)
assert 'num1' not in res_df.columns
assert 'num2' in res_df.columns
assert 'char' in res_df.columns
def test_coldrop_1(self):
"""Testing the ColDrop pipline stage."""
df = pd.DataFrame(
data=[[1, 'a'], [2, 'b']],
index=[1, 2],
columns=['num', 'char']
)
self.assertEqual(
collections.Counter(df.columns),
collections.Counter(['num', 'char']))
self.assertTrue('num' in df.columns)
stage = ColDrop('num')
df = stage.apply(df)
self.assertTrue('num' not in df.columns)
self.assertTrue('char' in df.columns)
def test_coldrop_non_str_lbl():
"""Testing the ColDrop pipeline stage."""
df = _test_df2()
assert 2 in df.columns
stage = ColDrop(2)
res_df = stage.apply(df)
assert 2 not in res_df.columns
assert 'num1' in res_df.columns
assert False in res_df.columns
def test_coldrop_missing_col():
"""Testing the ColDrop pipeline stage."""
df = _test_df()
assert 'num1' in df.columns
stage = ColDrop('num3')
with pytest.raises(FailedPreconditionError):
res_df = stage.apply(df)
stage = ColDrop('num3', errors='ignore')
res_df = stage.apply(df)
assert res_df.equals(df)
def test_coldrop_missing_col():
"""Testing the ColDrop pipeline stage."""
df = _test_df()
assert 'num1' in df.columns
stage = ColDrop('num3')
with pytest.raises(FailedPreconditionError):
res_df = stage.apply(df)
stage = ColDrop('num3', errors='ignore')
res_df = stage.apply(df)
assert res_df.equals(df)
def test_coldrop_multi_col():
"""Testing the ColDrop pipeline stage."""
df = _test_df()
assert 'num1' in df.columns
assert 'num2' in df.columns
stage = ColDrop(['num1', 'num2'])
res_df = stage.apply(df)
assert 'num1' not in res_df.columns
assert 'num2' not in res_df.columns
assert 'char' in res_df.columns
def test_attribute_stage():
"""Testing attribute pipeline stages."""
pipeline = pdp.ColDrop('name').Bin({'speed': [5]}, drop=True)
assert isinstance(pipeline, PdPipeline)
assert isinstance(pipeline[0], ColDrop)
assert isinstance(pipeline[1], Bin)
df = _some_df()
res_df = pipeline(df)
assert 'speed' in res_df.columns
assert 'name' not in res_df.columns
def test_coldrop_lambda():
"""Testing the ColDrop pipeline stage."""
df = _test_df()
assert 'num1' in df.columns
assert 'num2' in df.columns
stage = ColDrop(lambda col: 'num' in col.name)
res_df = stage.apply(df)
assert 'num1' not in res_df.columns
assert 'num2' not in res_df.columns
assert 'char' in res_df.columns
def __init__(self, columns, errors=None, **kwargs):
self._columns = columns
self._errors = errors
self._columns_str = _list_str(self._columns)
if not callable(columns):
self._columns = _interpret_columns_param(columns)
super_kwargs = {
'exmsg': ColDrop._DEF_COLDROP_EXC_MSG.format(self._columns_str),
'appmsg': ColDrop._DEF_COLDROP_APPLY_MSG.format(self._columns_str),
'desc': self._default_desc()
}
super_kwargs.update(**kwargs)
super().__init__(**super_kwargs)
def __init__(self, columns, errors=None, **kwargs):
self._columns = columns
self._errors = errors
self._columns_str = _list_str(self._columns)
if not callable(columns):
self._columns = _interpret_columns_param(columns)
super_kwargs = {
'exmsg': ColDrop._DEF_COLDROP_EXC_MSG.format(self._columns_str),
'appmsg': ColDrop._DEF_COLDROP_APPLY_MSG.format(self._columns_str),
'desc': self._default_desc()
}
super_kwargs.update(**kwargs)
super().__init__(**super_kwargs)