[ENH] defer trailing underscore attribute assignment in fit() for imputers and discretisers

feature_engine/_base_transformers/base_numerical.py

@@ -28,18 +28,18 @@ class BaseNumericalTransformer(
     variable transformers, discretisers, math combination.
     """
 
-    def fit(self, X: pd.DataFrame) -> pd.DataFrame:
+    def _fit_setup(self, X: pd.DataFrame):
         """
         Checks that input is a dataframe, finds numerical variables, or alternatively
-        checks that variables entered by the user are of type numerical.
+        checks that variables entered by the user are of type numerical, and checks
+        for NA and Inf. Does not assign any trailing-underscore attribute, so that
+        subclasses can defer attribute assignment until the rest of their fit logic
+        has completed successfully.
 
         Parameters
         ----------
         X : Pandas DataFrame
 
-        y : Pandas Series, np.array. Default = None
-            Parameter is necessary for compatibility with sklearn Pipeline.
-
         Raises
         ------
         TypeError
@@ -53,30 +53,34 @@ def fit(self, X: pd.DataFrame) -> pd.DataFrame:
         -------
         X : Pandas DataFrame
             The same dataframe entered as parameter
+
+        variables_ : List
+            The variables that were found or checked.
         """
 
         # check input dataframe
         X = check_X(X)
 
         # find or check for numerical variables
         if self.variables is None:
-            self.variables_ = find_numerical_variables(
-                X, return_empty=self.return_empty
-            )
+            variables_ = find_numerical_variables(X, return_empty=self.return_empty)
         else:
-            self.variables_ = check_numerical_variables(X, self.variables)
+            variables_ = check_numerical_variables(X, self.variables)
 
         # check if dataset contains na or inf
-        _check_contains_na(X, self.variables_)
-        _check_contains_inf(X, self.variables_)
+        _check_contains_na(X, variables_)
+        _check_contains_inf(X, variables_)
 
-        # save input features
-        self.feature_names_in_ = X.columns.tolist()
+        return X, variables_
+
+    def _get_feature_names_in(self, X):
+        """Get the names and number of features in the train set (the dataframe
+        used during fit)."""
 
-        # save train set shape
+        self.feature_names_in_ = X.columns.tolist()
         self.n_features_in_ = X.shape[1]
 
-        return X
+        return self
 
     def _check_transform_input_and_state(self, X: pd.DataFrame) -> pd.DataFrame:
         """

feature_engine/_base_transformers/mixins.py

@@ -1,4 +1,4 @@
-from typing import Dict, List, Union
+from typing import Dict, List, Tuple, Union
 
 import pandas as pd
 from numpy import ndarray
@@ -46,10 +46,14 @@ def transform_x_y(self, X: pd.DataFrame, y: pd.Series):
 
 
 class FitFromDictMixin:
-    def _fit_from_dict(self, X: pd.DataFrame, user_dict_: Dict) -> pd.DataFrame:
+    def _fit_from_dict(
+        self, X: pd.DataFrame, user_dict_: Dict
+    ) -> Tuple[pd.DataFrame, List[Union[str, int]]]:
         """
         Checks that input is a dataframe, checks that variables in the dictionary
-        entered by the user are of type numerical.
+        entered by the user are of type numerical. Does not assign any
+        trailing-underscore attribute, so that subclasses can defer attribute
+        assignment until the rest of their fit logic has completed successfully.
 
         Parameters
         ----------
@@ -71,25 +75,22 @@ def _fit_from_dict(self, X: pd.DataFrame, user_dict_: Dict) -> pd.DataFrame:
         -------
         X : Pandas DataFrame
             The same dataframe entered as parameter
+
+        variables_ : List
+            The variables in the dictionary.
         """
         # check input dataframe
         X = check_X(X)
 
         # find or check for numerical variables
         variables = list(user_dict_.keys())
-        self.variables_ = check_numerical_variables(X, variables)
+        variables_ = check_numerical_variables(X, variables)
 
         # check if dataset contains na or inf
-        _check_contains_na(X, self.variables_)
-        _check_contains_inf(X, self.variables_)
-
-        # save input features
-        self.feature_names_in_ = X.columns.tolist()
-
-        # save train set shape
-        self.n_features_in_ = X.shape[1]
+        _check_contains_na(X, variables_)
+        _check_contains_inf(X, variables_)
 
-        return X
+        return X, variables_
 
 
 class GetFeatureNamesOutMixin:

feature_engine/creation/cyclical_features.py

@@ -155,11 +155,15 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None):
             It is not needed in this transformer. You can pass y or None.
         """
         if self.max_values is None:
-            X = super().fit(X)
-            self.max_values_ = X[self.variables_].max().to_dict()
+            X, variables_ = self._fit_setup(X)
+            max_values_ = X[variables_].max().to_dict()
         else:
-            super()._fit_from_dict(X, self.max_values)
-            self.max_values_ = self.max_values
+            X, variables_ = super()._fit_from_dict(X, self.max_values)
+            max_values_ = self.max_values
+
+        self.variables_ = variables_
+        self.max_values_ = max_values_
+        self._get_feature_names_in(X)
 
         return self
 

feature_engine/creation/geo_features.py

@@ -234,26 +234,26 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None):
         # check input dataframe
         X = check_X(X)
 
-        # Store coordinate variables
-        self.variables_: List[Union[str, int]] = [
+        # Coordinate variables
+        variables: List[Union[str, int]] = [
             self.lat1,
             self.lon1,
             self.lat2,
             self.lon2,
         ]
 
         # Check all coordinate columns exist
-        missing = set(self.variables_) - set(X.columns)
+        missing = set(variables) - set(X.columns)
         if missing:
             raise ValueError(
                 f"Coordinate columns {missing} are not present in the dataframe."
             )
 
         # Check coordinate columns are numerical
-        check_numerical_variables(X, self.variables_)
+        check_numerical_variables(X, variables)
 
         # Check for missing values
-        _check_contains_na(X, self.variables_)
+        _check_contains_na(X, variables)
 
         # Validate coordinate ranges if enabled
         if self.validate_ranges:
@@ -269,6 +269,9 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None):
                         f"Longitude values in '{lon_col}' must be between -180 and 180."
                     )
 
+        # save coordinate variables
+        self.variables_ = variables
+
         # save input features
         self.feature_names_in_ = X.columns.tolist()
 

feature_engine/discretisation/arbitrary.py

@@ -152,10 +152,12 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None):
             y is not needed in this transformer. You can pass y or None.
         """
         # check input dataframe
-        X = super()._fit_from_dict(X, self.binning_dict)
+        X, variables_ = super()._fit_from_dict(X, self.binning_dict)
 
+        self.variables_ = variables_
         # for consistency with the rest of the discretisers, we add this attribute
         self.binner_dict_ = self.binning_dict
+        self._get_feature_names_in(X)
 
         return self
 

feature_engine/discretisation/base_discretiser.py

@@ -10,7 +10,8 @@ class BaseDiscretiser(BaseNumericalTransformer):
     """
     Shared set-up checks and methods across numerical discretisers.
 
-    Important: inherits fit() functionality and tags from BaseNumericalTransformer.
+    Important: inherits _fit_setup(), _get_feature_names_in() and tags from
+    BaseNumericalTransformer. Subclasses implement fit() themselves.
     """
 
     def __init__(

feature_engine/discretisation/decision_tree.py

@@ -225,7 +225,7 @@ def __init__(
         self.random_state = random_state
         self.return_empty = return_empty
 
-    def fit(self, X: pd.DataFrame, y: pd.Series):  # type: ignore
+    def fit(self, X: pd.DataFrame, y: pd.Series):
         """
         Fit one decision tree per variable to discretise with cross-validation and
         grid-search for hyperparameters.
@@ -252,7 +252,7 @@ def fit(self, X: pd.DataFrame, y: pd.Series):  # type: ignore
             check_classification_targets(y)
 
         # check input dataframe
-        X = super().fit(X)
+        X, variables_ = self._fit_setup(X)
 
         if self.param_grid:
             param_grid = self.param_grid
@@ -262,7 +262,7 @@ def fit(self, X: pd.DataFrame, y: pd.Series):  # type: ignore
         binner_dict_ = {}
         scores_dict_ = {}
 
-        for var in self.variables_:
+        for var in variables_:
 
             if self.regression:
                 model = DecisionTreeRegressor(random_state=self.random_state)
@@ -280,7 +280,7 @@ def fit(self, X: pd.DataFrame, y: pd.Series):  # type: ignore
             scores_dict_[var] = tree_model.score(X[var].to_frame(), y)
 
         if self.bin_output != "prediction":
-            for var in self.variables_:
+            for var in variables_:
                 clf = binner_dict_[var].best_estimator_
                 threshold = clf.tree_.threshold
                 feature = clf.tree_.feature
@@ -291,6 +291,9 @@ def fit(self, X: pd.DataFrame, y: pd.Series):  # type: ignore
 
         self.binner_dict_ = binner_dict_
         self.scores_dict_ = scores_dict_
+        self.variables_ = variables_
+        self._get_feature_names_in(X)
+
         return self
 
     def transform(self, X: pd.DataFrame) -> pd.DataFrame:

feature_engine/discretisation/equal_frequency.py

@@ -170,17 +170,21 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None):
         """
 
         # check input dataframe
-        X = super().fit(X)
+        X, variables_ = self._fit_setup(X)
 
-        self.binner_dict_ = {}
+        binner_dict_ = {}
 
-        for var in self.variables_:
+        for var in variables_:
             tmp, bins = pd.qcut(x=X[var], q=self.q, retbins=True, duplicates="drop")
 
             # Prepend/Append infinities to accommodate outliers
             bins = list(bins)
             bins[0] = float("-inf")
             bins[len(bins) - 1] = float("inf")
-            self.binner_dict_[var] = bins
+            binner_dict_[var] = bins
+
+        self.binner_dict_ = binner_dict_
+        self.variables_ = variables_
+        self._get_feature_names_in(X)
 
         return self

feature_engine/discretisation/equal_width.py

@@ -179,12 +179,12 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None):
         """
 
         # check input dataframe
-        X = super().fit(X)
+        X, variables_ = self._fit_setup(X)
 
         # fit
-        self.binner_dict_ = {}
+        binner_dict_ = {}
 
-        for var in self.variables_:
+        for var in variables_:
             tmp, bins = pd.cut(
                 x=X[var],
                 bins=self.bins,
@@ -197,6 +197,10 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None):
             bins = list(bins)
             bins[0] = float("-inf")
             bins[len(bins) - 1] = float("inf")
-            self.binner_dict_[var] = bins
+            binner_dict_[var] = bins
+
+        self.binner_dict_ = binner_dict_
+        self.variables_ = variables_
+        self._get_feature_names_in(X)
 
         return self

feature_engine/discretisation/geometric_width.py

@@ -174,19 +174,23 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None):
         """
 
         # check input dataframe
-        X = super().fit(X)
+        X, variables_ = self._fit_setup(X)
 
         # fit
-        self.binner_dict_ = {}
+        binner_dict_ = {}
 
-        for var in self.variables_:
+        for var in variables_:
             min_, max_ = X[var].min(), X[var].max()
             increment = np.power(max_ - min_, 1.0 / self.bins)
             bins = np.r_[
                 -np.inf, min_ + np.power(increment, np.arange(1, self.bins)), np.inf
             ]
             bins = np.sort(bins)
             bins = list(bins)
-            self.binner_dict_[var] = bins
+            binner_dict_[var] = bins
+
+        self.binner_dict_ = binner_dict_
+        self.variables_ = variables_
+        self._get_feature_names_in(X)
 
         return self

feature_engine/imputation/arbitrary_imputer.py

@@ -163,17 +163,19 @@ def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None):
         # find or check for numerical variables
         # create the imputer dictionary
         if self.imputer_dict:
-            self.variables_ = check_numerical_variables(
+            variables_ = check_numerical_variables(
                 X, list(self.imputer_dict.keys())
             )
-            self.imputer_dict_ = self.imputer_dict
+            imputer_dict_ = self.imputer_dict
         else:
             if self.variables is None:
-                self.variables_ = find_numerical_variables(X, self.return_empty)
+                variables_ = find_numerical_variables(X, self.return_empty)
             else:
-                self.variables_ = check_numerical_variables(X, self.variables)
-            self.imputer_dict_ = {var: self.arbitrary_number for var in self.variables_}
+                variables_ = check_numerical_variables(X, self.variables)
+            imputer_dict_ = {var: self.arbitrary_number for var in variables_}
 
+        self.variables_ = variables_
+        self.imputer_dict_ = imputer_dict_
         self._get_feature_names_in(X)
 
         return self