Write custom transformer in sklearn which returns .predict of estimator in .transform

Question

We have a custom transformer

class EstimatorTransformer(base.BaseEstimator, base.TransformerMixin):

    def __init__(self, estimator):
        self.estimator = estimator

    def fit(self, X, y):
        self = self.estimator.fit(X,y)
        return self

    def transform(self, X):
        return self.estimator.predict(X)

And there is an assert statement

city_trans = EstimatorTransformer(city_est)
city_trans.fit(features,target)
assert ([r[0] for r in city_trans.transform(data[:5])]
        == city_est.predict(data[:5]))

where

city_est is the estimator we can pass. I am using city_est = city_est = Ridge(alpha = 1)

but I get an error in self = self.estimator.fit(X,y). What I might be doing wrong here. I know that fit() returns self. How should I make this assertion work?

Answer 1

You are doing wrong assignment in this line:

self = self.estimator.fit(X,y)

Here, self is the current class (EstimatorTransformer) and you are trying to assign it a different class.

You can just write:

def fit(self, X, y):
    self.estimator.fit(X,y)
    return self

and it will work.