from sklearn.base import BaseEstimator, TransformerMixin


class DropNaNRowsTransformer(BaseEstimator, TransformerMixin):
    """
    Drop rows containing missing values (NaNs) inside an sklearn pipeline.

    Parameters
    ----------
    features : list[str] | None, optional
        If provided, drop rows that contain NaN in any of the listed columns.
        If None, drop rows that contain NaN in any column.
    """
    def __init__(self, features=None):
        """
        features: list of feature names or None.
                  - If list: rows with NaN in any of these features will be dropped.
                  - If None: rows with NaN in any column will be dropped.
        """
        self.features = features

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

    def transform(self, X):
        X = X.copy()

        if self.features is None:
            return X.dropna()

        missing_cols = set(self.features) - set(X.columns)
        if missing_cols:
            raise ValueError(
                f"The following columns are not present in the DataFrame: {missing_cols}"
            )

        return X.dropna(subset=self.features)
