from funcy import rcompose


class PredictionModelHandle:
    """Simplifies usage of ModelHandle instances for prediction purposes."""

    def __init__(self, model_handle):
        self.__model_handle = model_handle

    def predict(self, *args, **kwargs):
        predict = rcompose(self.__model_handle.prep_images, self.__model_handle.model.predict)
        return predict(*args, **kwargs)

    def predict_proba(self, *args, **kwargs):
        predict = rcompose(self.__model_handle.prep_images, self.__model_handle.model.predict_proba)
        return predict(*args, **kwargs)