[FIXED] What type is a sklearn model?

Issue

I'm writing some code which evaluates different sklearn models against some data. I am using type hints, both for my own education and to help other people who will eventually have to read my code.

My question is how do I specify the type of a sklearn predictor (such as LinearRegression())?

For example:

def model_tester(model : Predictor,
                 parameter: int
                 ) -> np.ndarray:
     """An example function with type hints."""

     # do stuff to model 

     return values

I see the typing library can make new types or I can use TypeVar to do:

Predictor = TypeVar('Predictor') 

but I wouldn't want to use this if there was already a conventional type for an sklearn model.

Checking the type of LinearRegression() yields:

 sklearn.linear_model.base.LinearRegression

and this is clearly of use, but only if I am interested in the LinearRegression model.

Solution

I think the most generic class that all models inherit from would be sklearn.base.BaseEstimator.

If you want to be more specific, maybe use sklearn.base.ClassifierMixin or sklearn.base.RegressorMixin.

So I would do:

from sklearn.base import RegressorMixin


def model_tester(model: RegressorMixin, parameter: int) -> np.ndarray:
     """An example function with type hints."""

     # do stuff to model 

     return values

I am no expert in type checking, so correct me if this is not right.

Answered By - FlorianGD