diff --git a/feature_engine/imputation/arbitrary_number.py b/feature_engine/imputation/arbitrary_number.py index ead2c0c4d..16c367aa6 100644 --- a/feature_engine/imputation/arbitrary_number.py +++ b/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): """ diff --git a/feature_engine/outliers/artbitrary.py b/feature_engine/outliers/artbitrary.py index e6d316621..59d7c1709 100644 --- a/feature_engine/outliers/artbitrary.py +++ b/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): diff --git a/feature_engine/parameter_checks.py b/feature_engine/parameter_checks.py index 10d4b014b..3c4573f8b 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 diff --git a/tests/test_parameter_checks.py b/tests/test_parameter_checks.py index 6cb148931..e8e878b75 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)