Example of Scikit-Learn pipeline and estimator

estimator.py

class DataFrameSelector(BaseEstimator, TransformerMixin):
    
    def __init__(self, dtype=None, drop_columns=None):
        self.dtype = dtype
        self.columns = drop_columns
    
    def fit(self, X, y=None):
        return self
    
    def transform(self, X):
        return (
            X
            .pipe(lambda d: d.drop(columns=self.drop_columns) if self.drop_columns is not None else d)
            .pipe(lambda d: d.select_dtypes(include=['number']) if self.dtype == 'numerical' else d)
            .pipe(lambda d: d.select_dtypes(include=['object']) if self.dtype == 'categorical' else d)
        )
        
t = DataFrameSelector()
t.transform(X_train)

sklearn.py

from sklearn.model_selection import train_test_split
from sklearn.pipeline import Pipeline, FeatureUnion
from sklearn.impute import SimpleImputer
from sklearn.preprocessing import OneHotEncoder
from utilities import FeaturesDrop
from sklearn.ensemble import RandomForestClassifier

target = 'Survived'
X = titanic.drop(columns=[target])
y = titanic[target]

X_train, X_test, y_train, y_test = train_test_split(X, y, train_size=0.8)

num_pipeline = Pipeline([
    ('select_numerical', DataFrameSelector(dtype='numerical')),
    ('impute', SimpleImputer())
])

cat_pipeline = Pipeline([
    ('select_categorical', DataFrameSelector(dtype='categorical')),
    ('encode', OneHotEncoder()),
])

pipeline = Pipeline([
    ('drop_features', FeaturesDrop(columns=['PassengerId', 'Name'])),
    ('union', FeatureUnion(
        transformer_list=[
            ('num_pipeline', num_pipeline),
            ('cat_pipeline', cat_pipeline)
        ]
    ))
])

pipeline.fit(X_train, y_train)