Как объединить создания нескольких переменных в одном классе

Question

У меня есть набор данных, как показано ниже

import pandas as pd 
import sklearn 
df= pd.DataFrame({'color': ['red', 'red,blue','red,blue,yellow', 'red,yellow', 'blue,yellow']}) 

я получаю новую переменную как этот

df['red'] = 0 
df.ix[df['color'].str.contains("red") == True, 'red' ] =1 

аналогично я могу получить df['blue'] & df['yellow'] Тогда я должен был использовать это в class (я хочу применить pipeline)

class Red(): 

def transform(self, X, y=None, **fit_params): 
    X['red'] = 0 
    X.loc[X['color'].str.contains("red") == True, 'red' ] = 1 
    return X['red'].values.reshape(len(X),1) 

def fit_transform(self, X, y=None, **fit_params): 
    self.fit(X, y, **fit_params) 
    return self.transform(X) 

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

Это работает, но я хочу получить class, который также генерирует 'blue' & 'yellow'. Сделать классы для каждого «цвета»? В реальном наборе данных есть несколько десятков цветов ». Я новичок в классе, пожалуйста, дайте какие-либо идеи, как объединить в одном class порождающих нескольких переменных

Answer 1

Я был удивлен, но это работает

class Red(BaseEstimator, TransformerMixin): 

def transform(self, X, y=None, **fit_params): 
    X['red'] = 0 
    X.loc[X['color'].str.contains("red") == True, 'red' ] = 1 
    X['blue'] = 0 
    X.loc[X['color'].str.contains("blue") == True, 'blue' ] = 1 
    return X[['red', 'blue']].values.reshape(len(X),2) 

def fit_transform(self, X, y=None, **fit_params): 
    self.fit(X, y, **fit_params) 
    return self.transform(X) 

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