# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/sklegos.ipynb (unless otherwise specified).

__all__ = ['train_test_split', 'MakeFrame', 'ColumnSelector']

# Cell
from .loader import logger

try:
    from sklearn.model_selection import train_test_split
except ModuleNotFoundError:
    logger.Exception('Please `pip install scikit-learn sklego` to use this submodule')

train_test_split = train_test_split

# Cell

from sklearn.base import BaseEstimator, TransformerMixin, MetaEstimatorMixin
class MakeFrame(BaseEstimator, TransformerMixin):
    '''Convert sklearn's output to a pandas dataframe
    Especially useful when working with an ensemble of models
    '''
    def __init__(self, column_names):
        self.column_names = column_names

    def fit(self, X, y=None):
        return X

    def transform(self, X, y=None):
        return pd.DataFrame(X, columns=self.column_names)

from sklego.preprocessing import ColumnSelector
ColumnSelector = ColumnSelector