File size: 342 Bytes
0e7829a
 
 
533e5f2
 
0e7829a
 
 
 
 
 
533e5f2
1
2
3
4
5
6
7
8
9
10
11
12
from sklearn.base import BaseEstimator, TransformerMixin

class NumericSelector(BaseEstimator, TransformerMixin):
    def __init__(self):
        pass

    def fit(self, X, y=None):
        self.numeric_cols_ = X.select_dtypes(include=[float, int]).columns
        return self

    def transform(self, X):
        return X[self.numeric_cols_]