refactor numerical dictionary check

feature_engine/imputation/arbitrary_number.py

@@ -17,7 +17,7 @@
 )
 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 @@ def __init__(
         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):
         """

feature_engine/outliers/artbitrary.py

@@ -23,7 +23,7 @@
 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 @@ def __init__(
         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):

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

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)