From Python 3.8 on (or earlier using typing-extensions), you can use typing.Protocol
. Using protocols, you can use a concept called structural subtyping to define exactly the type’s expected structure:
from typing import Protocol
# from typing_extensions import Protocol # for Python <3.8
class ScikitModel(Protocol):
def fit(self, X, y, sample_weight=None): ...
def predict(self, X): ...
def score(self, X, y, sample_weight=None): ...
def set_params(self, **params): ...
which you can then use as a type hint:
def do_stuff(model: ScikitModel) -> Any:
model.fit(train_data, train_labels) # this type checks
score = model.score(test_data, test_labels) # this type checks
...