def _test_get_column_map_features_name(self, col_format, expected):
sk = SklearnTransformerMock()
sk.get_feature_names = MagicMock(return_value=["0_blue", "0_red"])
if col_format is not None:
t = r.SklearnColumnsWrapper(sk, col_format=col_format)
else:
t = r.SklearnColumnsWrapper(sk)
t.cols = ["color"]
d = t.get_column_mapping()
d_expected = {"color": tuple(expected)}
self.assertDictEqual(d, d_expected)
def test_fit_transform_features_name_result(self):
sk = SklearnTransformerMock()
sk.get_feature_names = MagicMock(return_value=["0_blue", "0_red"])
return_value = np.array([[1, 2, 3], [3, 2, 1]]).T
sk.transform = MagicMock(return_value=return_value)
t = r.SklearnColumnsWrapper(sk, "color")
df_out = t.fit_transform(create_df_all())
df_expected = pd.DataFrame({
"color=blue": [1, 2, 3],
"color=red": [3, 2, 1],
"price": [1.5, 2.5, 3.5],
"amount": [1, 2, 3]
})
sk.transform.assert_called_once()
self.assertTrue(df_out.equals(df_expected))
def test_get_column_map_features_name_mix(self):
"""One transformer with get_feature_names and other without. """
sk1 = SklearnTransformerMock()
return_value = np.array(["0_blue", "0_red"])
sk1.get_feature_names = MagicMock(return_value=return_value)
sk2 = SklearnTransformerMock()
t = r.SklearnColumnWrapper(None)
t.transformers = {"color": sk1, "amount": sk2}
t.cols = ["color", "amount"]
d = t.get_column_mapping()
d_expected = {
"color": tuple(["color=blue", "color=red"]),
"amount": "amount"
}
self.assertDictEqual(d, d_expected)