Repository: feature-engine/feature_engine
Pull request #506 (2022-08-18 11:23:52)

refactor numerical dictionary check


Patch:
diff --git a/feature_engine/dataframe_checks.py b/feature_engine/dataframe_checks.py
index a53a1c1..8269727 100644
--- a/feature_engine/dataframe_checks.py
+++ b/feature_engine/dataframe_checks.py
@@ -2,7 +2,7 @@
 transform().
 """
 
-from typing import List, Union, Tuple
+from typing import List, Tuple, Union
 
 import numpy as np
 import pandas as pd
@@ -191,7 +191,7 @@ def check_X_y(
     if isinstance(X, pd.DataFrame) and isinstance(y, pd.Series):
         X, y = _check_X_y(X, y)
         # Check that their indexes match.
-        if not all(y.index == X.index):
+        if X.index.equals(y.index) is False:
             raise ValueError("The indexes of X and y do not match.")
 
     # case 2: X is dataframe and y is something else
@@ -258,7 +258,7 @@ def _check_contains_na(X: pd.DataFrame, variables: List[Union[str, int]]) -> Non
         If the variable(s) contain null values
     """
 
-    if X[variables].isnull().values.any():
+    if X[variables].isnull().any().any():
         raise ValueError(
             "Some of the variables to transform contain NaN. Check and "
             "remove those before using this transformer."
@@ -281,7 +281,7 @@ def _check_contains_inf(X: pd.DataFrame, variables: List[Union[str, int]]) -> No
         If the variable(s) contain np.inf values
     """
 
-    if np.isinf(X[variables]).values.any():
+    if np.isinf(X[variables]).any().any():
         raise ValueError(
             "Some of the variables to transform contain inf values. Check and "
             "remove those before using this transformer."
diff --git a/feature_engine/get_feature_names_out.py b/feature_engine/get_feature_names_out.py
index 03440c7..620172f 100644
--- a/feature_engine/get_feature_names_out.py
+++ b/feature_engine/get_feature_names_out.py
@@ -1,15 +1,16 @@
-"""Method functionality shared by all transformers that do not add features to the data,
-instead modify features in place."""
+"""Method shared by all transformers that do not add features to the data."""
 
 from typing import List, Optional, Union
 
 
 def _get_feature_names_out(
-        features_in: List[Union[str, int]],
-        transformed_features: List[Union[str, int]],
-        input_features: Optional[List] = None,
+    features_in: List[Union[str, int]],
+    transformed_features: List[Union[str, int]],
+    input_features: Optional[List] = None,
 ) -> List[Union[str, int]]:
     """
+    Returns variable names of transformed dataframe.
+
     If input_features is None, returns the names of all variables in the transformed
     dataframe. If input_features is a list, returns the feature names in the list. This
     parameter exists mostly for compatibility with the Scikit-learn Pipeline.
@@ -46,9 +47,13 @@ def _get_feature_names_out(
             raise ValueError(
                 f"input_features must be a list. Got {input_features} instead."
             )
-        if any([
-            feature for feature in input_features if feature not in transformed_features
-        ]):
+        if any(
+            [
+                feature
+                for feature in input_features
+                if feature not in transformed_features
+            ]
+        ):
             raise ValueError(
                 "Some features in input_features were not transformed by this "
                 "transformer. Pass either None, or a list with the features "
diff --git a/feature_engine/imputation/arbitrary_number.py b/feature_engine/imputation/arbitrary_number.py
index ead2c0c..16c367a 100644
--- a/feature_engine/imputation/arbitrary_number.py
+++ b/feature_engine/imputation/arbitrary_number.py
@@ -17,7 +17,7 @@ from feature_engine._docstrings.fit_attributes import (
 )
 from feature_engine._docstrings.substitute import Substitution
 from feature_engine.imputation.base_imputer import BaseImputer
-from feature_engine.parameter_checks import _define_numerical_dict
+from feature_engine.parameter_checks import _check_numerical_dict
 from feature_engine.variable_manipulation import (
     _check_input_parameter_variables,
     _find_or_check_numerical_variables,
@@ -97,9 +97,11 @@ class ArbitraryNumberImputer(BaseImputer):
         else:
             raise ValueError("arbitrary_number must be numeric of type int or float")
 
+        _check_numerical_dict(imputer_dict)
+
         self.variables = _check_input_parameter_variables(variables)
 
-        self.imputer_dict = _define_numerical_dict(imputer_dict)
+        self.imputer_dict = imputer_dict
 
     def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None):
         """
diff --git a/feature_engine/outliers/artbitrary.py b/feature_engine/outliers/artbitrary.py
index e6d3166..59d7c17 100644
--- a/feature_engine/outliers/artbitrary.py
+++ b/feature_engine/outliers/artbitrary.py
@@ -23,7 +23,7 @@ from feature_engine._docstrings.fit_attributes import (
 from feature_engine._docstrings.class_inputs import _missing_values_docstring
 from feature_engine._docstrings.substitute import Substitution
 from feature_engine.outliers.base_outlier import BaseOutlier
-from feature_engine.parameter_checks import _define_numerical_dict
+from feature_engine.parameter_checks import _check_numerical_dict
 from feature_engine.tags import _return_tags
 from feature_engine.variable_manipulation import _find_or_check_numerical_variables
 
@@ -100,8 +100,11 @@ class ArbitraryOutlierCapper(BaseOutlier):
         if missing_values not in ["raise", "ignore"]:
             raise ValueError("missing_values takes only values 'raise' or 'ignore'")
 
-        self.max_capping_dict = _define_numerical_dict(max_capping_dict)
-        self.min_capping_dict = _define_numerical_dict(min_capping_dict)
+        _check_numerical_dict(max_capping_dict)
+        _check_numerical_dict(min_capping_dict)
+
+        self.max_capping_dict = max_capping_dict
+        self.min_capping_dict = min_capping_dict
         self.missing_values = missing_values
 
     def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None):
diff --git a/feature_engine/parameter_checks.py b/feature_engine/parameter_checks.py
index 10d4b01..3c4573f 100644
--- a/feature_engine/parameter_checks.py
+++ b/feature_engine/parameter_checks.py
@@ -1,7 +1,7 @@
 from typing import Optional
 
 
-def _define_numerical_dict(dict_: Optional[dict]) -> Optional[dict]:
+def _check_numerical_dict(dict_: Optional[dict]) -> Optional[dict]:
     """
     Checks if all values in dictionary are integers and floats. Can take None as
     argument.
@@ -17,20 +17,13 @@ def _define_numerical_dict(dict_: Optional[dict]) -> Optional[dict]:
         If any of the values in the dictionary are not int or float
     TypeError
         When argument type is not a dictionary.
-
-    Returns
-    -------
-    None or the input dictionary
     """
 
-    if not dict_:
-        dict_ = dict_
-
-    elif isinstance(dict_, dict):
+    if isinstance(dict_, dict):
         if not all([isinstance(x, (float, int)) for x in dict_.values()]):
             raise ValueError("All values in the dictionary must be integer or float")
 
-    else:
+    elif dict_ is not None:
         raise TypeError("The parameter can only take a dictionary or None")
 
-    return dict_
+    return None


Test patch:
diff --git a/tests/test_parameter_checks.py b/tests/test_parameter_checks.py
index 6cb1489..e8e878b 100644
--- a/tests/test_parameter_checks.py
+++ b/tests/test_parameter_checks.py
@@ -1,24 +1,15 @@
 import pytest
 
-from feature_engine.parameter_checks import _define_numerical_dict
+from feature_engine.parameter_checks import _check_numerical_dict
 
 
-def test_numerical_dict():
-    input_dict = {"a": 1, "b": 2}
-    expected_output = {"a": 1, "b": 2}
-
-    assert _define_numerical_dict(input_dict) == expected_output
-
-
-def test_not_numerical_dict():
-    input_dict = {"a": 1, "b": "c"}
-
+@pytest.mark.parametrize("input_dict", [{"a": 1, "b": "c"}, {1: 1, 2: "c"}])
+def test_not_numerical_dict(input_dict):
     with pytest.raises(ValueError):
-        assert _define_numerical_dict(input_dict)
-
+        _check_numerical_dict(input_dict)
 
-def test_input_type():
-    input_dict = [1, 2, 3]
 
+@pytest.mark.parametrize("input_dict", [[1, 2, 3], (1, 2, 3), "hola", 5])
+def test_input_type(input_dict):
     with pytest.raises(TypeError):
-        assert _define_numerical_dict(input_dict)
+        _check_numerical_dict(input_dict)