Repository: feature-engine/feature_engine

Issue:
no numerical or categorical variable checks when user passes string instead of list
`_find_or_check_numerical_variable`s and `_find_or_check_categorical_variables` don't check variable type when user passes string instead of list

Copying @dodoarg code that spots this issue:


Please have a look at the demo below (on the PowerTransformer, which I haven't touched in my fork) <br>
Here, _variables_ is passed as a list and the correct TypeError is thrown at the _find_or_check_numerical_variables level:
```
>>> pt = PowerTransformer(variables=["City"])
>>> pt.fit(df)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "D:\feature_engine\feature_engine\transformation\power.py", line 82, in fit
    X = super().fit(X)
  File "D:\feature_engine\feature_engine\base_transformers.py", line 98, in fit
    self.variables_ = _find_or_check_numerical_variables(X, self.variables)
  File "D:\feature_engine\feature_engine\variable_manipulation.py", line 82, in _find_or_check_numerical_variables
    raise TypeError(
TypeError: Some of the variables are not numerical. Please cast them as numerical before using this transformer
```
Here instead variables is passed as a str (which _check_input_parameter_variables __allows__) but _find_or_check_numerical_variables doesn't actually run any checks, so the logic goes on up until the line where the column is checked for the presence of inf, which fails because City is an object variable, and should not have made it that far in the code
```
>>> pt = PowerTransformer(variables="City")   
>>> pt.fit(df)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "D:\feature_engine\feature_engine\transformation\power.py", line 82, in fit
    X = super().fit(X)
  File "D:\feature_engine\feature_engine\base_transformers.py", line 102, in fit
    _check_contains_inf(X, self.variables_)
  File "D:\feature_engine\feature_engine\dataframe_checks.py", line 127, in _check_contains_inf
    if np.isinf(X[variables]).values.any():
  File "D:\Vari\Anaconda3\envs\ML\lib\site-packages\pandas\core\generic.py", line 1936, in __array_ufunc__
    return arraylike.array_ufunc(self, ufunc, method, *inputs, **kwargs)
  File "D:\Vari\Anaconda3\envs\ML\lib\site-packages\pandas\core\arraylike.py", line 366, in array_ufunc
    result = mgr.apply(getattr(ufunc, method))
  File "D:\Vari\Anaconda3\envs\ML\lib\site-packages\pandas\core\internals\managers.py", line 425, in apply
    applied = b.apply(f, **kwargs)
  File "D:\Vari\Anaconda3\envs\ML\lib\site-packages\pandas\core\internals\blocks.py", line 378, in apply
    result = func(self.values, **kwargs)
TypeError: ufunc 'isinf' not supported for the input types, and the inputs could not be safely coerced to any supported types according to the casting rule ''safe''
>>>
```


Discussion:
dodoarg: Sorry about the verbose responses @solegalli, I tried to keep it simple in the PR description but unfortunately this opened a small can of worms :/
solegalli: > Sorry about the verbose responses @solegalli, I tried to keep it simple in the PR description but unfortunately this opened a small can of worms :/

No problem at all, I am always happy to learn new stuff :) 

thanks for the detail and dedication
solegalli: Hi @dodoarg 

Suggested a few changes here:https://github.com/dodoarg/feature_engine/pull/2

This is ready to go. Thank you!
dodoarg: I believe this is ready to go now?
solegalli: @SunnyxBd small change here in one of the classes you are working on.

Patch:
diff --git a/feature_engine/selection/target_mean_selection.py b/feature_engine/selection/target_mean_selection.py
index c55717b..6eba350 100644
--- a/feature_engine/selection/target_mean_selection.py
+++ b/feature_engine/selection/target_mean_selection.py
@@ -201,7 +201,7 @@ class SelectByTargetMeanPerformance(BaseSelector):
         # find categorical and numerical variables
         self.variables_categorical_ = list(X.select_dtypes(include="O").columns)
         self.variables_numerical_ = list(
-            X.select_dtypes(include=["float", "integer"]).columns
+            X.select_dtypes(include="number").columns
         )
 
         # obtain cross-validation indeces
diff --git a/feature_engine/variable_manipulation.py b/feature_engine/variable_manipulation.py
index a8430cb..a043dad 100644
--- a/feature_engine/variable_manipulation.py
+++ b/feature_engine/variable_manipulation.py
@@ -4,6 +4,10 @@ from typing import Any, List, Union
 
 import pandas as pd
 
+from pandas.api.types import is_numeric_dtype as is_numeric
+from pandas.api.types import is_categorical_dtype as is_categorical
+from pandas.api.types import is_object_dtype as is_object
+
 Variables = Union[None, int, str, List[Union[str, int]]]
 
 
@@ -44,40 +48,47 @@ def _find_or_check_numerical_variables(
 
     Parameters
     ----------
-    X : Pandas DataFrame
+    X : Pandas DataFrame.
     variables : variable or list of variables. Defaults to None.
 
     Raises
     ------
     ValueError
-        If there are no numerical variables in the df or the df is empty
+        If there are no numerical variables in the df or the df is empty.
     TypeError
-        If any of the user provided variables are not numerical
+        If any of the user provided variables are not numerical.
 
     Returns
     -------
-    variables: List of numerical variables
+    variables: List of numerical variables.
     """
 
-    if isinstance(variables, (str, int)):
-        variables = [variables]
-
-    elif not variables:
+    if variables is None:
         # find numerical variables in dataset
         variables = list(X.select_dtypes(include="number").columns)
         if len(variables) == 0:
             raise ValueError(
-                "No numerical variables in this dataframe. Please check variable"
-                "format with pandas dtypes"
+                "No numerical variables found in this dataframe. Please check "
+                "variable format with pandas dtypes."
             )
 
+    elif isinstance(variables, (str, int)):
+        if is_numeric(X[variables]):
+            variables = [variables]
+        else:
+            raise TypeError("The variable entered is not numeric.")
+
     else:
+        if len(variables) == 0:
+            raise ValueError("The list of variables is empty.")
+
         # check that user entered variables are of type numerical
-        if any(X[variables].select_dtypes(exclude="number").columns):
-            raise TypeError(
-                "Some of the variables are not numerical. Please cast them as "
-                "numerical before using this transformer"
-            )
+        else:
+            if len(X[variables].select_dtypes(exclude="number").columns) > 0:
+                raise TypeError(
+                    "Some of the variables are not numerical. Please cast them as "
+                    "numerical before using this transformer."
+                )
 
     return variables
 
@@ -91,38 +102,47 @@ def _find_or_check_categorical_variables(
 
     Parameters
     ----------
-    X : pandas DataFrame
+    X : pandas DataFrame.
     variables : variable or list of variables. Defaults to None.
 
     Raises
     ------
     ValueError
-        If there are no categorical variables in df or df is empty
+        If there are no categorical variables in df or df is empty.
     TypeError
-        If any of the user provided variables are not categorical
+        If any of the user provided variables are not categorical.
 
     Returns
     -------
-    variables : List of categorical variables
+    variables : List of categorical variables.
     """
 
-    if isinstance(variables, (str, int)):
-        variables = [variables]
-
-    elif not variables:
+    if variables is None:
+        # find categorical variables in dataset
         variables = list(X.select_dtypes(include=["O", "category"]).columns)
         if len(variables) == 0:
             raise ValueError(
-                "No categorical variables in this dataframe. Please check the "
-                "variables format with pandas dtypes"
+                "No categorical variables found in this dataframe. Please check "
+                "variable format with pandas dtypes."
             )
 
+    elif isinstance(variables, (str, int)):
+        if is_categorical(X[variables]) or is_object(X[variables]):
+            variables = [variables]
+        else:
+            raise TypeError("The variable entered is not categorical.")
+
     else:
-        if any(X[variables].select_dtypes(exclude=["O", "category"]).columns):
-            raise TypeError(
-                "Some of the variables are not categorical. Please cast them as object "
-                "or category before calling this transformer"
-            )
+        if len(variables) == 0:
+            raise ValueError("The list of variables is empty.")
+
+        # check that user entered variables are of type numerical
+        else:
+            if len(X[variables].select_dtypes(exclude=["O", "category"]).columns) > 0:
+                raise TypeError(
+                    "Some of the variables are not categorical. Please cast them as "
+                    "categorical or object before using this transformer."
+                )
 
     return variables
 


Test patch:
diff --git a/tests/test_variable_manipulation.py b/tests/test_variable_manipulation.py
index 647eb70..eafb2e8 100644
--- a/tests/test_variable_manipulation.py
+++ b/tests/test_variable_manipulation.py
@@ -44,9 +44,21 @@ def test_find_or_check_numerical_variables(df_vartypes, df_numeric_columns):
     assert _find_or_check_numerical_variables(df_vartypes, var_num) == ["Age"]
     assert _find_or_check_numerical_variables(df_vartypes, vars_none) == vars_num
 
+    with pytest.raises(TypeError):
+        assert _find_or_check_numerical_variables(df_vartypes, "City")
+
+    with pytest.raises(TypeError):
+        assert _find_or_check_numerical_variables(df_numeric_columns, 0)
+
+    with pytest.raises(TypeError):
+        assert _find_or_check_numerical_variables(df_numeric_columns, [1, 3])
+
     with pytest.raises(TypeError):
         assert _find_or_check_numerical_variables(df_vartypes, vars_mix)
 
+    with pytest.raises(ValueError):
+        assert _find_or_check_numerical_variables(df_vartypes, variables=[])
+
     with pytest.raises(ValueError):
         assert _find_or_check_numerical_variables(df_vartypes[["Name", "City"]], None)
 
@@ -61,13 +73,26 @@ def test_find_or_check_categorical_variables(df_vartypes, df_numeric_columns):
     assert _find_or_check_categorical_variables(df_vartypes, vars_cat) == vars_cat
     assert _find_or_check_categorical_variables(df_vartypes, None) == vars_cat
 
+    with pytest.raises(TypeError):
+        assert _find_or_check_categorical_variables(df_vartypes, "Marks")
+
+    with pytest.raises(TypeError):
+        assert _find_or_check_categorical_variables(df_numeric_columns, 3)
+
+    with pytest.raises(TypeError):
+        assert _find_or_check_categorical_variables(df_numeric_columns, [0, 2])
+
     with pytest.raises(TypeError):
         assert _find_or_check_categorical_variables(df_vartypes, vars_mix)
 
+    with pytest.raises(ValueError):
+        assert _find_or_check_categorical_variables(df_vartypes, variables=[])
+
     with pytest.raises(ValueError):
         assert _find_or_check_categorical_variables(df_vartypes[["Age", "Marks"]], None)
 
     assert _find_or_check_categorical_variables(df_numeric_columns, [0, 1]) == [0, 1]
+    assert _find_or_check_categorical_variables(df_numeric_columns, 0) == [0]
     assert _find_or_check_categorical_variables(df_numeric_columns, 1) == [1]
 
     df_vartypes["Age"] = df_vartypes["Age"].astype("category")