venv/lib/python3.11/site-packages/sklearn/tests/metadata_routing_common.py

import inspect
from collections import defaultdict
from functools import partial

import numpy as np
from numpy.testing import assert_array_equal

from sklearn.base import (
    BaseEstimator,
    ClassifierMixin,
    MetaEstimatorMixin,
    RegressorMixin,
    TransformerMixin,
    clone,
)
from sklearn.metrics._scorer import _Scorer, mean_squared_error
from sklearn.model_selection import BaseCrossValidator
from sklearn.model_selection._split import GroupsConsumerMixin
from sklearn.utils._metadata_requests import (
    SIMPLE_METHODS,
)
from sklearn.utils.metadata_routing import (
    MetadataRouter,
    MethodMapping,
    process_routing,
)
from sklearn.utils.multiclass import _check_partial_fit_first_call


def record_metadata(obj, record_default=True, **kwargs):
    """Utility function to store passed metadata to a method of obj.

    If record_default is False, kwargs whose values are "default" are skipped.
    This is so that checks on keyword arguments whose default was not changed
    are skipped.

    """
    stack = inspect.stack()
    callee = stack[1].function
    caller = stack[2].function
    if not hasattr(obj, "_records"):
        obj._records = defaultdict(lambda: defaultdict(list))
    if not record_default:
        kwargs = {
            key: val
            for key, val in kwargs.items()
            if not isinstance(val, str) or (val != "default")
        }
    obj._records[callee][caller].append(kwargs)


def check_recorded_metadata(obj, method, parent, split_params=tuple(), **kwargs):
    """Check whether the expected metadata is passed to the object's method.

    Parameters
    ----------
    obj : estimator object
        sub-estimator to check routed params for
    method : str
        sub-estimator's method where metadata is routed to, or otherwise in
        the context of metadata routing referred to as 'callee'
    parent : str
        the parent method which should have called `method`, or otherwise in
        the context of metadata routing referred to as 'caller'
    split_params : tuple, default=empty
        specifies any parameters which are to be checked as being a subset
        of the original values
    **kwargs : dict
        passed metadata
    """
    all_records = (
        getattr(obj, "_records", dict()).get(method, dict()).get(parent, list())
    )
    for record in all_records:
        # first check that the names of the metadata passed are the same as
        # expected. The names are stored as keys in `record`.
        assert set(kwargs.keys()) == set(
            record.keys()
        ), f"Expected {kwargs.keys()} vs {record.keys()}"
        for key, value in kwargs.items():
            recorded_value = record[key]
            # The following condition is used to check for any specified parameters
            # being a subset of the original values
            if key in split_params and recorded_value is not None:
                assert np.isin(recorded_value, value).all()
            else:
                if isinstance(recorded_value, np.ndarray):
                    assert_array_equal(recorded_value, value)
                else:
                    assert (
                        recorded_value is value
                    ), f"Expected {recorded_value} vs {value}. Method: {method}"


record_metadata_not_default = partial(record_metadata, record_default=False)


def assert_request_is_empty(metadata_request, exclude=None):
    """Check if a metadata request dict is empty.

    One can exclude a method or a list of methods from the check using the
    ``exclude`` parameter. If metadata_request is a MetadataRouter, then
    ``exclude`` can be of the form ``{"object" : [method, ...]}``.
    """
    if isinstance(metadata_request, MetadataRouter):
        for name, route_mapping in metadata_request:
            if exclude is not None and name in exclude:
                _exclude = exclude[name]
            else:
                _exclude = None
            assert_request_is_empty(route_mapping.router, exclude=_exclude)
        return

    exclude = [] if exclude is None else exclude
    for method in SIMPLE_METHODS:
        if method in exclude:
            continue
        mmr = getattr(metadata_request, method)
        props = [
            prop
            for prop, alias in mmr.requests.items()
            if isinstance(alias, str) or alias is not None
        ]
        assert not props


def assert_request_equal(request, dictionary):
    for method, requests in dictionary.items():
        mmr = getattr(request, method)
        assert mmr.requests == requests

    empty_methods = [method for method in SIMPLE_METHODS if method not in dictionary]
    for method in empty_methods:
        assert not len(getattr(request, method).requests)


class _Registry(list):
    # This list is used to get a reference to the sub-estimators, which are not
    # necessarily stored on the metaestimator. We need to override __deepcopy__
    # because the sub-estimators are probably cloned, which would result in a
    # new copy of the list, but we need copy and deep copy both to return the
    # same instance.
    def __deepcopy__(self, memo):
        return self

    def __copy__(self):
        return self


class ConsumingRegressor(RegressorMixin, BaseEstimator):
    """A regressor consuming metadata.

    Parameters
    ----------
    registry : list, default=None
        If a list, the estimator will append itself to the list in order to have
        a reference to the estimator later on. Since that reference is not
        required in all tests, registration can be skipped by leaving this value
        as None.
    """

    def __init__(self, registry=None):
        self.registry = registry

    def partial_fit(self, X, y, sample_weight="default", metadata="default"):
        if self.registry is not None:
            self.registry.append(self)

        record_metadata_not_default(
            self, sample_weight=sample_weight, metadata=metadata
        )
        return self

    def fit(self, X, y, sample_weight="default", metadata="default"):
        if self.registry is not None:
            self.registry.append(self)

        record_metadata_not_default(
            self, sample_weight=sample_weight, metadata=metadata
        )
        return self

    def predict(self, X, y=None, sample_weight="default", metadata="default"):
        record_metadata_not_default(
            self, sample_weight=sample_weight, metadata=metadata
        )
        return np.zeros(shape=(len(X),))

    def score(self, X, y, sample_weight="default", metadata="default"):
        record_metadata_not_default(
            self, sample_weight=sample_weight, metadata=metadata
        )
        return 1


class NonConsumingClassifier(ClassifierMixin, BaseEstimator):
    """A classifier which accepts no metadata on any method."""

    def __init__(self, alpha=0.0):
        self.alpha = alpha

    def fit(self, X, y):
        self.classes_ = np.unique(y)
        self.coef_ = np.ones_like(X)
        return self

    def partial_fit(self, X, y, classes=None):
        return self

    def decision_function(self, X):
        return self.predict(X)

    def predict(self, X):
        y_pred = np.empty(shape=(len(X),))
        y_pred[: len(X) // 2] = 0
        y_pred[len(X) // 2 :] = 1
        return y_pred

    def predict_proba(self, X):
        # dummy probabilities to support predict_proba
        y_proba = np.empty(shape=(len(X), 2))
        y_proba[: len(X) // 2, :] = np.asarray([1.0, 0.0])
        y_proba[len(X) // 2 :, :] = np.asarray([0.0, 1.0])
        return y_proba

    def predict_log_proba(self, X):
        # dummy probabilities to support predict_log_proba
        return self.predict_proba(X)


class NonConsumingRegressor(RegressorMixin, BaseEstimator):
    """A classifier which accepts no metadata on any method."""

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

    def partial_fit(self, X, y):
        return self

    def predict(self, X):
        return np.ones(len(X))  # pragma: no cover


class ConsumingClassifier(ClassifierMixin, BaseEstimator):
    """A classifier consuming metadata.

    Parameters
    ----------
    registry : list, default=None
        If a list, the estimator will append itself to the list in order to have
        a reference to the estimator later on. Since that reference is not
        required in all tests, registration can be skipped by leaving this value
        as None.

    alpha : float, default=0
        This parameter is only used to test the ``*SearchCV`` objects, and
        doesn't do anything.
    """

    def __init__(self, registry=None, alpha=0.0):
        self.alpha = alpha
        self.registry = registry

    def partial_fit(
        self, X, y, classes=None, sample_weight="default", metadata="default"
    ):
        if self.registry is not None:
            self.registry.append(self)

        record_metadata_not_default(
            self, sample_weight=sample_weight, metadata=metadata
        )
        _check_partial_fit_first_call(self, classes)
        return self

    def fit(self, X, y, sample_weight="default", metadata="default"):
        if self.registry is not None:
            self.registry.append(self)

        record_metadata_not_default(
            self, sample_weight=sample_weight, metadata=metadata
        )

        self.classes_ = np.unique(y)
        self.coef_ = np.ones_like(X)
        return self

    def predict(self, X, sample_weight="default", metadata="default"):
        record_metadata_not_default(
            self, sample_weight=sample_weight, metadata=metadata
        )
        y_score = np.empty(shape=(len(X),), dtype="int8")
        y_score[len(X) // 2 :] = 0
        y_score[: len(X) // 2] = 1
        return y_score

    def predict_proba(self, X, sample_weight="default", metadata="default"):
        record_metadata_not_default(
            self, sample_weight=sample_weight, metadata=metadata
        )
        y_proba = np.empty(shape=(len(X), 2))
        y_proba[: len(X) // 2, :] = np.asarray([1.0, 0.0])
        y_proba[len(X) // 2 :, :] = np.asarray([0.0, 1.0])
        return y_proba

    def predict_log_proba(self, X, sample_weight="default", metadata="default"):
        record_metadata_not_default(
            self, sample_weight=sample_weight, metadata=metadata
        )
        return np.zeros(shape=(len(X), 2))

    def decision_function(self, X, sample_weight="default", metadata="default"):
        record_metadata_not_default(
            self, sample_weight=sample_weight, metadata=metadata
        )
        y_score = np.empty(shape=(len(X),))
        y_score[len(X) // 2 :] = 0
        y_score[: len(X) // 2] = 1
        return y_score

    def score(self, X, y, sample_weight="default", metadata="default"):
        record_metadata_not_default(
            self, sample_weight=sample_weight, metadata=metadata
        )
        return 1


class ConsumingTransformer(TransformerMixin, BaseEstimator):
    """A transformer which accepts metadata on fit and transform.

    Parameters
    ----------
    registry : list, default=None
        If a list, the estimator will append itself to the list in order to have
        a reference to the estimator later on. Since that reference is not
        required in all tests, registration can be skipped by leaving this value
        as None.
    """

    def __init__(self, registry=None):
        self.registry = registry

    def fit(self, X, y=None, sample_weight="default", metadata="default"):
        if self.registry is not None:
            self.registry.append(self)

        record_metadata_not_default(
            self, sample_weight=sample_weight, metadata=metadata
        )
        self.fitted_ = True
        return self

    def transform(self, X, sample_weight="default", metadata="default"):
        record_metadata_not_default(
            self, sample_weight=sample_weight, metadata=metadata
        )
        return X + 1

    def fit_transform(self, X, y, sample_weight="default", metadata="default"):
        # implementing ``fit_transform`` is necessary since
        # ``TransformerMixin.fit_transform`` doesn't route any metadata to
        # ``transform``, while here we want ``transform`` to receive
        # ``sample_weight`` and ``metadata``.
        record_metadata_not_default(
            self, sample_weight=sample_weight, metadata=metadata
        )
        return self.fit(X, y, sample_weight=sample_weight, metadata=metadata).transform(
            X, sample_weight=sample_weight, metadata=metadata
        )

    def inverse_transform(self, X, sample_weight=None, metadata=None):
        record_metadata_not_default(
            self, sample_weight=sample_weight, metadata=metadata
        )
        return X - 1


class ConsumingNoFitTransformTransformer(BaseEstimator):
    """A metadata consuming transformer that doesn't inherit from
    TransformerMixin, and thus doesn't implement `fit_transform`. Note that
    TransformerMixin's `fit_transform` doesn't route metadata to `transform`."""

    def __init__(self, registry=None):
        self.registry = registry

    def fit(self, X, y=None, sample_weight=None, metadata=None):
        if self.registry is not None:
            self.registry.append(self)

        record_metadata(self, sample_weight=sample_weight, metadata=metadata)

        return self

    def transform(self, X, sample_weight=None, metadata=None):
        record_metadata(self, sample_weight=sample_weight, metadata=metadata)
        return X


class ConsumingScorer(_Scorer):
    def __init__(self, registry=None):
        super().__init__(
            score_func=mean_squared_error, sign=1, kwargs={}, response_method="predict"
        )
        self.registry = registry

    def _score(self, method_caller, clf, X, y, **kwargs):
        if self.registry is not None:
            self.registry.append(self)

        record_metadata_not_default(self, **kwargs)

        sample_weight = kwargs.get("sample_weight", None)
        return super()._score(method_caller, clf, X, y, sample_weight=sample_weight)


class ConsumingSplitter(GroupsConsumerMixin, BaseCrossValidator):
    def __init__(self, registry=None):
        self.registry = registry

    def split(self, X, y=None, groups="default", metadata="default"):
        if self.registry is not None:
            self.registry.append(self)

        record_metadata_not_default(self, groups=groups, metadata=metadata)

        split_index = len(X) // 2
        train_indices = list(range(0, split_index))
        test_indices = list(range(split_index, len(X)))
        yield test_indices, train_indices
        yield train_indices, test_indices

    def get_n_splits(self, X=None, y=None, groups=None, metadata=None):
        return 2

    def _iter_test_indices(self, X=None, y=None, groups=None):
        split_index = len(X) // 2
        train_indices = list(range(0, split_index))
        test_indices = list(range(split_index, len(X)))
        yield test_indices
        yield train_indices


class MetaRegressor(MetaEstimatorMixin, RegressorMixin, BaseEstimator):
    """A meta-regressor which is only a router."""

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

    def fit(self, X, y, **fit_params):
        params = process_routing(self, "fit", **fit_params)
        self.estimator_ = clone(self.estimator).fit(X, y, **params.estimator.fit)

    def get_metadata_routing(self):
        router = MetadataRouter(owner=self.__class__.__name__).add(
            estimator=self.estimator,
            method_mapping=MethodMapping().add(caller="fit", callee="fit"),
        )
        return router


class WeightedMetaRegressor(MetaEstimatorMixin, RegressorMixin, BaseEstimator):
    """A meta-regressor which is also a consumer."""

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

    def fit(self, X, y, sample_weight=None, **fit_params):
        if self.registry is not None:
            self.registry.append(self)

        record_metadata(self, sample_weight=sample_weight)
        params = process_routing(self, "fit", sample_weight=sample_weight, **fit_params)
        self.estimator_ = clone(self.estimator).fit(X, y, **params.estimator.fit)
        return self

    def predict(self, X, **predict_params):
        params = process_routing(self, "predict", **predict_params)
        return self.estimator_.predict(X, **params.estimator.predict)

    def get_metadata_routing(self):
        router = (
            MetadataRouter(owner=self.__class__.__name__)
            .add_self_request(self)
            .add(
                estimator=self.estimator,
                method_mapping=MethodMapping()
                .add(caller="fit", callee="fit")
                .add(caller="predict", callee="predict"),
            )
        )
        return router


class WeightedMetaClassifier(MetaEstimatorMixin, ClassifierMixin, BaseEstimator):
    """A meta-estimator which also consumes sample_weight itself in ``fit``."""

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

    def fit(self, X, y, sample_weight=None, **kwargs):
        if self.registry is not None:
            self.registry.append(self)

        record_metadata(self, sample_weight=sample_weight)
        params = process_routing(self, "fit", sample_weight=sample_weight, **kwargs)
        self.estimator_ = clone(self.estimator).fit(X, y, **params.estimator.fit)
        return self

    def get_metadata_routing(self):
        router = (
            MetadataRouter(owner=self.__class__.__name__)
            .add_self_request(self)
            .add(
                estimator=self.estimator,
                method_mapping=MethodMapping().add(caller="fit", callee="fit"),
            )
        )
        return router


class MetaTransformer(MetaEstimatorMixin, TransformerMixin, BaseEstimator):
    """A simple meta-transformer."""

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

    def fit(self, X, y=None, **fit_params):
        params = process_routing(self, "fit", **fit_params)
        self.transformer_ = clone(self.transformer).fit(X, y, **params.transformer.fit)
        return self

    def transform(self, X, y=None, **transform_params):
        params = process_routing(self, "transform", **transform_params)
        return self.transformer_.transform(X, **params.transformer.transform)

    def get_metadata_routing(self):
        return MetadataRouter(owner=self.__class__.__name__).add(
            transformer=self.transformer,
            method_mapping=MethodMapping()
            .add(caller="fit", callee="fit")
            .add(caller="transform", callee="transform"),
        )