From cfe46453d1e31247cbf97f2f50a3e5c800767ecd Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Tue, 25 Aug 2026 01:42:56 +0200 Subject: [PATCH 1/5] Migrate CategoricalMethodsMixin (encoding base) to narwhals, add polars support Shared base for all 8 encoders. _get_feature_names_in() and _check_transform_input_and_state() follow the same is_pandas-gated column-reorder pattern as BaseImputer/DecisionTreeFeatures. _check_or_select_variables() needed no change: the variable_handling helpers it calls are already fully narwhals-generic. The hot path is _encode()/inverse_transform(), a per-column dict-based map applied on every transform() call across every encoder. Benchmarked pandas-native .map(dict) vs narwhals Series.replace_strict(dict, default=...) at 10k/50k/100k rows x 1/2/10 columns x 5/50 categories (warmed up first to remove first-call JIT/import overhead): narwhals-on-pandas lands at ~1.06x-1.2x of pandas-native at realistic sizes (50k-100k rows), i.e. minimal loss - merged into a single narwhals path per the established decision rule, no pandas fast-path split. narwhals-on- polars is consistently ~4-5x faster than pandas-native at 100k rows. replace_strict() also *simplifies* the old logic: pandas' plain .map() leaves category-dtype columns as category dtype after mapping, which the old code corrected with a manual "cast to int if all-int else float" step. Verified narwhals' replace_strict resolves straight to a plain numeric dtype on both a pandas category column and a polars Categorical column, so that dtype fixup is dead code once replace_strict replaces .map() - dropped it entirely rather than porting it. Used Series.get_column().replace_strict() (not nw.col(), which only accepts string names) throughout, same as DecisionTreeFeatures' precedent for pandas integer column names - nw.col(feature) blew up on int-named columns (caught by the existing test_column_names_are_numbers test, which polars can't cover since it has no integer-column-name concept). _check_nan_values_after_transformation() rewritten off pandas' .isnull().sum().sum()/.columns[...] chain onto per-column Series.null_count(), for the same int-column-name reason. Verified: tests/test_encoding full suite unchanged (17 pre-existing failures - numpy-array-input rejection per the narwhals check_X() contract, plus 3 MeanEncoder inverse_transform failures caused by a pre-existing bug in mean_encoding.py's still-unmigrated fit() passing a numpy y into y.groupby(); reproduced identically against the unmodified base_encoder.py to confirm neither predates nor is introduced by this change - 326 passed both before and after, same failing test IDs). flake8 and mypy clean on the file. Module imports with pandas blocked (loaded standalone, since sibling encoder files in this package are not yet migrated and still import pandas at their own module level). sphinx -W build clean (only the pre-existing unrelated linkcode_resolve warning). Manually verified CountEncoder end-to-end on polars input (fit still pandas-only until its own migration, transform/inverse_transform now backend-agnostic via this mixin) produces identical values to the pandas path, including a pre-existing quirk where count-encoding inverse_transform is ambiguous for categories that share a count (confirmed identical, not a regression, on the old code too). _helper_functions.py checked: pure-python parameter validation, no dataframe interaction, no pandas import - left untouched. Co-Authored-By: Claude Sonnet 5 --- feature_engine/encoding/base_encoder.py | 115 ++++++++++++++---------- 1 file changed, 69 insertions(+), 46 deletions(-) diff --git a/feature_engine/encoding/base_encoder.py b/feature_engine/encoding/base_encoder.py index 53eca3095..e2c4395f1 100644 --- a/feature_engine/encoding/base_encoder.py +++ b/feature_engine/encoding/base_encoder.py @@ -1,7 +1,9 @@ import warnings from typing import List, Union -import pandas as pd +import narwhals as nw +import narwhals.dependencies as nwd +from narwhals.typing import IntoDataFrame from sklearn.base import BaseEstimator, TransformerMixin from sklearn.utils.validation import check_is_fitted @@ -121,11 +123,11 @@ class CategoricalMethodsMixin(TransformerMixin, BaseEstimator, GetFeatureNamesOu - GetFeatureNamesOutMixin brings method get_feature_names_out(). """ - def _check_na(self, X: pd.DataFrame, variables): + def _check_na(self, X: IntoDataFrame, variables): if self.missing_values == "raise": _check_contains_na(X, variables, error_msg="optional") - def _check_or_select_variables(self, X: pd.DataFrame): + def _check_or_select_variables(self, X: IntoDataFrame): """ Finds categorical variables, or alternatively checks that the variables entered by the user are of type object (categorical). @@ -133,7 +135,7 @@ def _check_or_select_variables(self, X: pd.DataFrame): Parameters ---------- - X: Pandas DataFrame + X: dataframe Raises ------ @@ -159,37 +161,41 @@ def _check_or_select_variables(self, X: pd.DataFrame): return variables_ - def _get_feature_names_in(self, X: pd.DataFrame): + def _get_feature_names_in(self, X: IntoDataFrame): """ Returns attributes `featrure_names_in_` and `n_feature_names_in_`, which are standard for all transformers in the library. """ # save input features - self.feature_names_in_ = X.columns.tolist() + is_pandas = nwd.is_pandas_dataframe(X) + if is_pandas is True: + self.feature_names_in_ = list(X.columns) + else: + self.feature_names_in_ = nw.from_native(X, eager_only=True).columns # save train set shape self.n_features_in_ = X.shape[1] - def _check_transform_input_and_state(self, X: pd.DataFrame) -> pd.DataFrame: + def _check_transform_input_and_state(self, X: IntoDataFrame) -> IntoDataFrame: """ Checks that the input is a dataframe and of the same size than the one used in the fit method. Checks absence of NA. Parameters ---------- - X: Pandas DataFrame + X: dataframe Raises ------ TypeError - If the input is not a Pandas DataFrame + If the input is not a dataframe ValueError - If the variable(s) contain null values. - If the df has different number of features than the df used in fit() Returns ------- - X: Pandas DataFrame + X: dataframe The same dataframe entered by the user. """ @@ -203,21 +209,29 @@ def _check_transform_input_and_state(self, X: pd.DataFrame) -> pd.DataFrame: _check_X_matches_training_df(X, self.n_features_in_) # reorder df to match train set - X = X[self.feature_names_in_] + is_pandas = nwd.is_pandas_dataframe(X) + if is_pandas is True: + X = X[self.feature_names_in_] + else: + X = ( + nw.from_native(X, eager_only=True) + .select(self.feature_names_in_) + .to_native() + ) return X - def transform(self, X: pd.DataFrame) -> pd.DataFrame: + def transform(self, X: IntoDataFrame) -> IntoDataFrame: """Replace categories with the learned parameters. Parameters ---------- - X: pandas dataframe of shape = [n_samples, n_features]. + X: dataframe of shape = [n_samples, n_features]. The dataset to transform. Returns ------- - X_new: pandas dataframe of shape = [n_samples, n_features]. + X_new: dataframe of shape = [n_samples, n_features]. The dataframe containing the categories replaced by numbers. """ @@ -231,22 +245,25 @@ def transform(self, X: pd.DataFrame) -> pd.DataFrame: return X - def _encode(self, X: pd.DataFrame) -> pd.DataFrame: - # replace categories by the learned parameters - for feature in self.encoder_dict_.keys(): - X[feature] = X[feature].map(self.encoder_dict_[feature]) - - # if original variables are cast as categorical, they will remain - # categorical after the encoding, and this is probably not desired - if X[feature].dtype.name == "category": - if all(isinstance(x, int) for x in X[feature]): - X[feature] = X[feature].astype("int") - else: - X[feature] = X[feature].astype("float") - - if self.unseen == "encode": - X[self.variables_] = X[self.variables_].fillna(self._unseen) - else: + def _encode(self, X: IntoDataFrame) -> IntoDataFrame: + # replace categories by the learned parameters. + # narwhals' replace_strict() lets one expression both map known + # categories and fill unseen/missing ones via `default`, so the + # pandas-only category-dtype fixup this used to need (map() leaves + # category dtype behind) is no longer necessary: replace_strict + # already resolves to a plain numeric dtype on both backends. + # get_column()/Series.replace_strict() (rather than nw.col(), which + # only accepts string names) is what lets this handle pandas + # integer column names too, same as DecisionTreeFeatures. + default = self._unseen if self.unseen == "encode" else None + nw_X = nw.from_native(X, eager_only=True) + new_series = [ + nw_X.get_column(feature).replace_strict(mapping, default=default) + for feature, mapping in self.encoder_dict_.items() + ] + X = nw_X.with_columns(*new_series).to_native() + + if self.unseen != "encode": # check if nan values were introduced by the transformation self._check_nan_values_after_transformation(X) @@ -255,19 +272,19 @@ def _encode(self, X: pd.DataFrame) -> pd.DataFrame: def _check_nan_values_after_transformation(self, X): # check if NaN values were introduced by the encoding - if X[self.variables_].isnull().sum().sum() > 0: + nw_X = nw.from_native(X, eager_only=True) + nan_columns = [ + feature + for feature in self.encoder_dict_.keys() + if nw_X.get_column(feature).null_count() > 0 + ] - # obtain the name(s) of the columns have null values - nan_columns = ( - X[self.encoder_dict_.keys()] - .columns[X[self.encoder_dict_.keys()].isnull().any()] - .tolist() - ) + if len(nan_columns) > 0: if len(nan_columns) > 1: - nan_columns_str = ", ".join(nan_columns) + nan_columns_str = ", ".join(str(col) for col in nan_columns) else: - nan_columns_str = nan_columns[0] + nan_columns_str = str(nan_columns[0]) if self.unseen == "ignore": warnings.warn( @@ -280,27 +297,33 @@ def _check_nan_values_after_transformation(self, X): f"{nan_columns_str}." ) - def inverse_transform(self, X: pd.DataFrame) -> pd.DataFrame: + def inverse_transform(self, X: IntoDataFrame) -> IntoDataFrame: """Convert the encoded variable back to the original values. Parameters ---------- - X: pandas dataframe of shape = [n_samples, n_features]. + X: dataframe of shape = [n_samples, n_features]. The transformed dataframe. Returns ------- - X_tr: pandas dataframe of shape = [n_samples, n_features]. + X_tr: dataframe of shape = [n_samples, n_features]. The un-transformed dataframe, with the categorical variables containing the original values. """ X = self._check_transform_input_and_state(X) - # replace encoded categories by the original values - for feature in self.encoder_dict_.keys(): - inv_map = {v: k for k, v in self.encoder_dict_[feature].items()} - X[feature] = X[feature].map(inv_map) + # replace encoded categories by the original values. get_column() + # rather than nw.col() again, to support pandas integer column names. + nw_X = nw.from_native(X, eager_only=True) + new_series = [ + nw_X.get_column(feature).replace_strict( + {v: k for k, v in mapping.items()}, default=None + ) + for feature, mapping in self.encoder_dict_.items() + ] + X = nw_X.with_columns(*new_series).to_native() return X From 5c3a55b58ac7fe0c9305bd0b45d7e17338306d26 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Sun, 30 Aug 2026 21:34:53 +0200 Subject: [PATCH 2/5] Update base_encoder.py --- feature_engine/encoding/base_encoder.py | 43 ++++++------------------- 1 file changed, 9 insertions(+), 34 deletions(-) diff --git a/feature_engine/encoding/base_encoder.py b/feature_engine/encoding/base_encoder.py index e2c4395f1..7a9dc4c96 100644 --- a/feature_engine/encoding/base_encoder.py +++ b/feature_engine/encoding/base_encoder.py @@ -167,11 +167,7 @@ def _get_feature_names_in(self, X: IntoDataFrame): standard for all transformers in the library. """ # save input features - is_pandas = nwd.is_pandas_dataframe(X) - if is_pandas is True: - self.feature_names_in_ = list(X.columns) - else: - self.feature_names_in_ = nw.from_native(X, eager_only=True).columns + self.feature_names_in_ = X.columns # save train set shape self.n_features_in_ = X.shape[1] @@ -203,23 +199,12 @@ def _check_transform_input_and_state(self, X: IntoDataFrame) -> IntoDataFrame: check_is_fitted(self) # check that input is a dataframe - X = check_X(X) + nw_X = check_X(X) # Check input data contains same number of columns as df used to fit _check_X_matches_training_df(X, self.n_features_in_) - # reorder df to match train set - is_pandas = nwd.is_pandas_dataframe(X) - if is_pandas is True: - X = X[self.feature_names_in_] - else: - X = ( - nw.from_native(X, eager_only=True) - .select(self.feature_names_in_) - .to_native() - ) - - return X + return nw_X def transform(self, X: IntoDataFrame) -> IntoDataFrame: """Replace categories with the learned parameters. @@ -235,44 +220,35 @@ def transform(self, X: IntoDataFrame) -> IntoDataFrame: The dataframe containing the categories replaced by numbers. """ - X = self._check_transform_input_and_state(X) + nw_X = self._check_transform_input_and_state(X) # check if dataset contains na if self.missing_values == "raise": _check_contains_na(X, self.variables_, error_msg="optional") - X = self._encode(X) + X = self._encode(nw_X) return X def _encode(self, X: IntoDataFrame) -> IntoDataFrame: - # replace categories by the learned parameters. - # narwhals' replace_strict() lets one expression both map known - # categories and fill unseen/missing ones via `default`, so the - # pandas-only category-dtype fixup this used to need (map() leaves - # category dtype behind) is no longer necessary: replace_strict - # already resolves to a plain numeric dtype on both backends. - # get_column()/Series.replace_strict() (rather than nw.col(), which - # only accepts string names) is what lets this handle pandas - # integer column names too, same as DecisionTreeFeatures. default = self._unseen if self.unseen == "encode" else None - nw_X = nw.from_native(X, eager_only=True) new_series = [ nw_X.get_column(feature).replace_strict(mapping, default=default) for feature, mapping in self.encoder_dict_.items() ] - X = nw_X.with_columns(*new_series).to_native() + X = nw_X.with_columns(*new_series) if self.unseen != "encode": # check if nan values were introduced by the transformation self._check_nan_values_after_transformation(X) + + X = X.to_native() return X def _check_nan_values_after_transformation(self, X): # check if NaN values were introduced by the encoding - nw_X = nw.from_native(X, eager_only=True) nan_columns = [ feature for feature in self.encoder_dict_.keys() @@ -312,11 +288,10 @@ def inverse_transform(self, X: IntoDataFrame) -> IntoDataFrame: original values. """ - X = self._check_transform_input_and_state(X) + nw_X = self._check_transform_input_and_state(X) # replace encoded categories by the original values. get_column() # rather than nw.col() again, to support pandas integer column names. - nw_X = nw.from_native(X, eager_only=True) new_series = [ nw_X.get_column(feature).replace_strict( {v: k for k, v in mapping.items()}, default=None From 8003f07e50b297ef427bad475bcecf9ef0428833 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Sun, 30 Aug 2026 22:11:59 +0200 Subject: [PATCH 3/5] Fix CategoricalMethodsMixin for narwhals-returning check_X After the rebase onto narwhals-migration, check_X / check_X_y return a narwhals frame. The previous "Update base_encoder.py" left the method bodies referencing a local nw_X that no longer exists. - _encode / _check_nan_values_after_transformation: use the narwhals frame that is actually passed in (was NameError on nw_X). - _check_nan_values_after_transformation now assumes a narwhals frame (its only caller, _encode, hands it one); no nw.from_native round-trip. - _get_feature_names_in: single branch-free `list(X.columns)` (normalises a narwhals column list and a pandas Index alike). - _check_transform_input_and_state keeps the native X for the column-count check and returns the narwhals frame. - Drop now-unused narwhals imports; refresh docstrings. - test_categorical_method_mixin: pass a narwhals frame to the two direct _check_nan_values_after_transformation calls. The encoder subclasses still run pandas-only fit()/transform() code and are adapted in their own migration PRs. Co-Authored-By: Claude Sonnet 5 --- feature_engine/encoding/base_encoder.py | 47 +++++++++++-------- .../test_categorical_method_mixin.py | 5 +- 2 files changed, 31 insertions(+), 21 deletions(-) diff --git a/feature_engine/encoding/base_encoder.py b/feature_engine/encoding/base_encoder.py index 7a9dc4c96..1e58d1222 100644 --- a/feature_engine/encoding/base_encoder.py +++ b/feature_engine/encoding/base_encoder.py @@ -1,8 +1,6 @@ import warnings from typing import List, Union -import narwhals as nw -import narwhals.dependencies as nwd from narwhals.typing import IntoDataFrame from sklearn.base import BaseEstimator, TransformerMixin from sklearn.utils.validation import check_is_fitted @@ -163,11 +161,17 @@ def _check_or_select_variables(self, X: IntoDataFrame): def _get_feature_names_in(self, X: IntoDataFrame): """ - Returns attributes `featrure_names_in_` and `n_feature_names_in_`, which are + Sets attributes `feature_names_in_` and `n_features_in_`, which are standard for all transformers in the library. + + Parameters + ---------- + X: narwhals dataframe + The dataframe returned by `check_X` / `check_X_y` at the start of `fit`. """ - # save input features - self.feature_names_in_ = X.columns + # save input features. list() normalises both a narwhals `.columns` + # (already a list) and a pandas `Index` to a plain list. + self.feature_names_in_ = list(X.columns) # save train set shape self.n_features_in_ = X.shape[1] @@ -175,30 +179,31 @@ def _get_feature_names_in(self, X: IntoDataFrame): def _check_transform_input_and_state(self, X: IntoDataFrame) -> IntoDataFrame: """ Checks that the input is a dataframe and of the same size than the one used - in the fit method. Checks absence of NA. + in the fit method. Parameters ---------- X: dataframe + The dataframe entered by the user, in any library supported by narwhals. Raises ------ TypeError If the input is not a dataframe ValueError - - If the variable(s) contain null values. - - If the df has different number of features than the df used in fit() + If the df has a different number of features than the df used in fit() Returns ------- - X: dataframe - The same dataframe entered by the user. + nw_X: narwhals dataframe + The narwhalified version of the dataframe entered by the user. """ # Check method fit has been called check_is_fitted(self) - # check that input is a dataframe + # check that input is a dataframe. check_X returns a narwhals frame; the + # original native X is kept for the column-count check below. nw_X = check_X(X) # Check input data contains same number of columns as df used to fit @@ -231,28 +236,32 @@ def transform(self, X: IntoDataFrame) -> IntoDataFrame: return X def _encode(self, X: IntoDataFrame) -> IntoDataFrame: + # X is the narwhals frame returned by _check_transform_input_and_state(). + # replace_strict() maps known categories and fills unseen/missing ones via + # `default` in a single expression, and resolves to a plain numeric dtype + # on both pandas and polars. get_column()/Series.replace_strict() (rather + # than nw.col(), which only accepts string names) is what lets this handle + # pandas integer column names too. default = self._unseen if self.unseen == "encode" else None new_series = [ - nw_X.get_column(feature).replace_strict(mapping, default=default) + X.get_column(feature).replace_strict(mapping, default=default) for feature, mapping in self.encoder_dict_.items() ] - X = nw_X.with_columns(*new_series) + X = X.with_columns(*new_series) if self.unseen != "encode": # check if nan values were introduced by the transformation self._check_nan_values_after_transformation(X) - - X = X.to_native() - - return X - def _check_nan_values_after_transformation(self, X): + return X.to_native() + def _check_nan_values_after_transformation(self, X: IntoDataFrame): + # X is the encoded narwhals frame built by _encode(). # check if NaN values were introduced by the encoding nan_columns = [ feature for feature in self.encoder_dict_.keys() - if nw_X.get_column(feature).null_count() > 0 + if X.get_column(feature).null_count() > 0 ] if len(nan_columns) > 0: diff --git a/tests/test_encoding/test_base_encoders/test_categorical_method_mixin.py b/tests/test_encoding/test_base_encoders/test_categorical_method_mixin.py index 80dc30813..ae28ce648 100644 --- a/tests/test_encoding/test_base_encoders/test_categorical_method_mixin.py +++ b/tests/test_encoding/test_base_encoders/test_categorical_method_mixin.py @@ -1,3 +1,4 @@ +import narwhals as nw import numpy as np import pandas as pd import pytest @@ -103,7 +104,7 @@ def test_raises_error_when_nan_introduced(): msg = "During the encoding, NaN values were introduced in the feature(s) words." with pytest.raises(ValueError) as record: - enc._check_nan_values_after_transformation(output_df) + enc._check_nan_values_after_transformation(nw.from_native(output_df)) assert str(record.value) == msg with pytest.raises(ValueError) as record: @@ -122,7 +123,7 @@ def test_raises_warning_when_nan_introduced(): assert record[0].message.args[0] == msg with pytest.warns(UserWarning) as record: - enc._check_nan_values_after_transformation(output_df) + enc._check_nan_values_after_transformation(nw.from_native(output_df)) assert record[0].message.args[0] == msg From 94f74f1561fe2475991aae2a4f39376c19c4a21a Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Sun, 30 Aug 2026 22:46:20 +0200 Subject: [PATCH 4/5] Update base_encoder.py --- feature_engine/encoding/base_encoder.py | 12 ------------ 1 file changed, 12 deletions(-) diff --git a/feature_engine/encoding/base_encoder.py b/feature_engine/encoding/base_encoder.py index 1e58d1222..adf856dbd 100644 --- a/feature_engine/encoding/base_encoder.py +++ b/feature_engine/encoding/base_encoder.py @@ -199,14 +199,10 @@ def _check_transform_input_and_state(self, X: IntoDataFrame) -> IntoDataFrame: The narwhalified version of the dataframe entered by the user. """ - # Check method fit has been called check_is_fitted(self) - # check that input is a dataframe. check_X returns a narwhals frame; the - # original native X is kept for the column-count check below. nw_X = check_X(X) - # Check input data contains same number of columns as df used to fit _check_X_matches_training_df(X, self.n_features_in_) return nw_X @@ -236,12 +232,6 @@ def transform(self, X: IntoDataFrame) -> IntoDataFrame: return X def _encode(self, X: IntoDataFrame) -> IntoDataFrame: - # X is the narwhals frame returned by _check_transform_input_and_state(). - # replace_strict() maps known categories and fills unseen/missing ones via - # `default` in a single expression, and resolves to a plain numeric dtype - # on both pandas and polars. get_column()/Series.replace_strict() (rather - # than nw.col(), which only accepts string names) is what lets this handle - # pandas integer column names too. default = self._unseen if self.unseen == "encode" else None new_series = [ X.get_column(feature).replace_strict(mapping, default=default) @@ -256,8 +246,6 @@ def _encode(self, X: IntoDataFrame) -> IntoDataFrame: return X.to_native() def _check_nan_values_after_transformation(self, X: IntoDataFrame): - # X is the encoded narwhals frame built by _encode(). - # check if NaN values were introduced by the encoding nan_columns = [ feature for feature in self.encoder_dict_.keys() From e5e5e08318a83418abd9360b5c86722d5290f758 Mon Sep 17 00:00:00 2001 From: Soledad Galli Date: Sun, 30 Aug 2026 22:49:28 +0200 Subject: [PATCH 5/5] test(encoding): assert error/warning text via pytest.raises/warns match= Replace the `as record: ... assert str(record.value) == msg` / `record[0].message.args[0] == msg` pattern in the CategoricalMethodsMixin tests with `match=re.escape(msg)` on pytest.raises / pytest.warns. Co-Authored-By: Claude Sonnet 5 --- .../test_categorical_method_mixin.py | 24 ++++++++----------- 1 file changed, 10 insertions(+), 14 deletions(-) diff --git a/tests/test_encoding/test_base_encoders/test_categorical_method_mixin.py b/tests/test_encoding/test_base_encoders/test_categorical_method_mixin.py index ae28ce648..609a4f9c4 100644 --- a/tests/test_encoding/test_base_encoders/test_categorical_method_mixin.py +++ b/tests/test_encoding/test_base_encoders/test_categorical_method_mixin.py @@ -1,3 +1,5 @@ +import re + import narwhals as nw import numpy as np import pandas as pd @@ -24,14 +26,13 @@ def test_underscore_check_na_method(): variables = ["words", "animals"] enc = MockClassFit(missing_values="raise") - with pytest.raises(ValueError) as record: - enc._check_na(input_df, variables) msg = ( "Some of the variables in the dataset contain NaN. Check and " "remove those before using this transformer or set the parameter " "`missing_values='ignore'` when initialising this transformer." ) - assert str(record.value) == msg + with pytest.raises(ValueError, match=re.escape(msg)): + enc._check_na(input_df, variables) def test_check_or_select_variables(): @@ -103,13 +104,11 @@ def test_raises_error_when_nan_introduced(): enc = MockClass(unseen="raise") msg = "During the encoding, NaN values were introduced in the feature(s) words." - with pytest.raises(ValueError) as record: + with pytest.raises(ValueError, match=re.escape(msg)): enc._check_nan_values_after_transformation(nw.from_native(output_df)) - assert str(record.value) == msg - with pytest.raises(ValueError) as record: + with pytest.raises(ValueError, match=re.escape(msg)): enc.transform(input_df) - assert str(record.value) == msg def test_raises_warning_when_nan_introduced(): @@ -118,26 +117,23 @@ def test_raises_warning_when_nan_introduced(): enc = MockClass(unseen="ignore") msg = "During the encoding, NaN values were introduced in the feature(s) words." - with pytest.warns(UserWarning) as record: + with pytest.warns(UserWarning, match=re.escape(msg)): enc.transform(input_df) - assert record[0].message.args[0] == msg - with pytest.warns(UserWarning) as record: + with pytest.warns(UserWarning, match=re.escape(msg)): enc._check_nan_values_after_transformation(nw.from_native(output_df)) - assert record[0].message.args[0] == msg def test_transform_raises_error_when_df_has_nan(): input_df = pd.DataFrame({"words": ["dog", "dig", "cat", np.nan]}) enc = MockClass() - with pytest.raises(ValueError) as record: - enc.transform(input_df) msg = ( "Some of the variables in the dataset contain NaN. Check and " "remove those before using this transformer or set the parameter " "`missing_values='ignore'` when initialising this transformer." ) - assert str(record.value) == msg + with pytest.raises(ValueError, match=re.escape(msg)): + enc.transform(input_df) def test_transform_ignores_nan_in_df_to_transform():