move function find_numerical_and_categorical_variables to find_variables.py

feature_engine/_prediction/base_predictor.py

@@ -22,9 +22,7 @@
 )
 from feature_engine.encoding import MeanEncoder
 from feature_engine.tags import _return_tags
-from feature_engine.variable_handling.variable_selection import (
-    find_categorical_and_numerical_variables,
-)
+from feature_engine.variable_handling import find_categorical_and_numerical_variables
 
 
 class BaseTargetMeanEstimator(BaseEstimator):

feature_engine/selection/drop_psi_features.py

@@ -33,9 +33,7 @@
 )
 from feature_engine.selection.base_selector import BaseSelector
 from feature_engine.tags import _return_tags
-from feature_engine.variable_handling.variable_selection import (
-    find_categorical_and_numerical_variables,
-)
+from feature_engine.variable_handling import find_categorical_and_numerical_variables
 
 Variables = Union[None, int, str, List[Union[str, int]]]
 

feature_engine/variable_handling/__init__.py

@@ -11,14 +11,13 @@
 )
 from .find_variables import (
     find_all_variables,
+    find_categorical_and_numerical_variables,
     find_categorical_variables,
     find_datetime_variables,
     find_numerical_variables,
 )
-from .variable_selection import find_categorical_and_numerical_variables
 
 __all__ = [
-    "find_categorical_and_numerical_variables",
     "check_all_variables",
     "check_numerical_variables",
     "check_categorical_variables",
@@ -27,4 +26,5 @@
     "find_numerical_variables",
     "find_categorical_variables",
     "find_datetime_variables",
+    "find_categorical_and_numerical_variables",
 ]

feature_engine/variable_handling/find_variables.py

@@ -1,7 +1,9 @@
-from typing import List, Union
+from typing import List, Tuple, Union
 
 import pandas as pd
 from pandas.api.types import is_datetime64_any_dtype as is_datetime
+from pandas.core.dtypes.common import is_numeric_dtype as is_numeric
+from pandas.core.dtypes.common import is_object_dtype as is_object
 
 from feature_engine.variable_handling._variable_type_checks import (
     _is_categorical_and_is_datetime,
@@ -178,3 +180,100 @@ def find_all_variables(
     else:
         variables = X.columns.to_list()
     return variables
+
+
+def find_categorical_and_numerical_variables(
+    X: pd.DataFrame,
+    variables: Union[None, int, str, List[Union[str, int]]] = None,
+) -> Tuple[List[Union[str, int]], List[Union[str, int]]]:
+    """
+    Find numerical and categorical variables in a dataframe or from a list.
+
+    The function returns two lists; the first one with the names of the variables of
+    type object or categorical and the second list with the names of the numerical
+    variables.
+
+    More details in the :ref:`User Guide <find_cat_and_num_vars>`.
+
+    Parameters
+    ----------
+    X : pandas dataframe of shape = [n_samples, n_features]
+        The dataset
+
+    variables : list, default=None
+        If `None`, the function will find all categorical and numerical variables in X.
+        Alternatively, it will find categorical and numerical variables in X, selecting
+        from the given list.
+
+    Returns
+    -------
+    variables: tuple
+        Tupe containing a list with the categorical variables, and a List with the
+        numerical variables.
+
+    Examples
+    --------
+    >>> import pandas as pd
+    >>> from feature_engine.variable_handling import (
+    >>>   find_categorical_and_numerical_variables
+    >>>)
+    >>> X = pd.DataFrame({
+    >>>     "var_num": [1, 2, 3],
+    >>>     "var_cat": ["A", "B", "C"],
+    >>>     "var_date": pd.date_range("2020-02-24", periods=3, freq="T")
+    >>> })
+    >>> var_cat, var_num = find_categorical_and_numerical_variables(X)
+    >>> var_cat, var_num
+    (['var_cat'], ['var_num'])
+    """
+
+    # If the user passes just 1 variable outside a list.
+    if isinstance(variables, (str, int)):
+
+        if X[variables].dtype.name == "category" or is_object(X[variables]):
+            variables_cat = [variables]
+            variables_num = []
+        elif is_numeric(X[variables]):
+            variables_num = [variables]
+            variables_cat = []
+        else:
+            raise TypeError(
+                "The variable entered is neither numerical nor categorical."
+            )
+
+    # If user leaves default None parameter.
+    elif variables is None:
+        # find categorical variables
+        if variables is None:
+            variables_cat = [
+                column
+                for column in X.select_dtypes(include=["O", "category"]).columns
+                if _is_categorical_and_is_not_datetime(X[column])
+            ]
+        # find numerical variables in dataset
+        variables_num = list(X.select_dtypes(include="number").columns)
+
+        if len(variables_num) == 0 and len(variables_cat) == 0:
+            raise TypeError(
+                "There are no numerical or categorical variables in the dataframe"
+            )
+
+    # If user passes variable list.
+    else:
+        if len(variables) == 0:
+            raise ValueError("The list of variables is empty.")
+
+        # find categorical variables
+        variables_cat = [
+            var for var in X[variables].select_dtypes(include=["O", "category"]).columns
+        ]
+
+        # find numerical variables
+        variables_num = list(X[variables].select_dtypes(include="number").columns)
+
+        if any([v for v in variables if v not in variables_cat + variables_num]):
+            raise TypeError(
+                "Some of the variables are neither numerical nor categorical."
+            )
+
+    return variables_cat, variables_num

feature_engine/variable_handling/variable_selection.py

@@ -1,14 +1,6 @@
 """Functions to select certain types of variables."""
 
-from typing import List, Tuple, Union
-
-import pandas as pd
-from pandas.api.types import is_numeric_dtype as is_numeric
-from pandas.api.types import is_object_dtype as is_object
-
-from feature_engine.variable_handling._variable_type_checks import (
-    _is_categorical_and_is_not_datetime,
-)
+from typing import List, Union
 
 Variables = Union[None, int, str, List[Union[str, int]]]
 
@@ -52,98 +44,3 @@ def _filter_out_variables_not_in_dataframe(X, variables):
         )
 
     return filtered_variables
-
-
-def find_categorical_and_numerical_variables(
-    X: pd.DataFrame, variables: Variables = None
-) -> Tuple[List[Union[str, int]], List[Union[str, int]]]:
-    """
-    Find numerical and categorical variables.
-
-    The function returns two lists; the first one with the names of the variables of
-    type object or categorical and the second list with the names of the numerical
-    variables.
-
-    More details in the :ref:`User Guide <find_cat_and_num_vars>`.
-
-    Parameters
-    ----------
-    X : pandas dataframe of shape = [n_samples, n_features]
-        The dataset
-
-    variables : list, default=None
-        If `None`, the function will find categorical and numerical variables in X.
-        Alternatively, it will find categorical and numerical variables in the given
-        list.
-
-    Returns
-    -------
-    variables: tuple
-        List of numerical and list of categorical variables.
-
-    Examples
-    --------
-    >>> import pandas as pd
-    >>> from feature_engine.variable_handling import (
-    >>>   find_categorical_and_numerical_variables
-    >>>)
-    >>> X = pd.DataFrame({
-    >>>     "var_num": [1, 2, 3],
-    >>>     "var_cat": ["A", "B", "C"],
-    >>>     "var_date": pd.date_range("2020-02-24", periods=3, freq="T")
-    >>> })
-    >>> var_cat, var_num = find_categorical_and_numerical_variables(X)
-    >>> var_cat, var_num
-    (['var_cat'], ['var_num'])
-    """
-
-    # If the user passes just 1 variable outside a list.
-    if isinstance(variables, (str, int)):
-
-        if X[variables].dtype.name == "category" or is_object(X[variables]):
-            variables_cat = [variables]
-            variables_num = []
-        elif is_numeric(X[variables]):
-            variables_num = [variables]
-            variables_cat = []
-        else:
-            raise TypeError(
-                "The variable entered is neither numerical nor categorical."
-            )
-
-    # If user leaves default None parameter.
-    elif variables is None:
-        # find categorical variables
-        if variables is None:
-            variables_cat = [
-                column
-                for column in X.select_dtypes(include=["O", "category"]).columns
-                if _is_categorical_and_is_not_datetime(X[column])
-            ]
-        # find numerical variables in dataset
-        variables_num = list(X.select_dtypes(include="number").columns)
-
-        if len(variables_num) == 0 and len(variables_cat) == 0:
-            raise TypeError(
-                "There are no numerical or categorical variables in the dataframe"
-            )
-
-    # If user passes variable list.
-    else:
-        if len(variables) == 0:
-            raise ValueError("The list of variables is empty.")
-
-        # find categorical variables
-        variables_cat = [
-            var for var in X[variables].select_dtypes(include=["O", "category"]).columns
-        ]
-
-        # find numerical variables
-        variables_num = list(X[variables].select_dtypes(include="number").columns)
-
-        if any([v for v in variables if v not in variables_cat + variables_num]):
-            raise TypeError(
-                "Some of the variables are neither numerical nor categorical."
-            )
-
-    return variables_cat, variables_num

tests/test_variable_handling/test_find_variables.py

@@ -2,6 +2,7 @@
 
 from feature_engine.variable_handling import (
     find_all_variables,
+    find_categorical_and_numerical_variables,
     find_categorical_variables,
     find_datetime_variables,
     find_numerical_variables,
@@ -80,3 +81,82 @@ def test_find_all_variables(df_vartypes):
 
     assert find_all_variables(df_vartypes, exclude_datetime=False) == all_vars
     assert find_all_variables(df_vartypes, exclude_datetime=True) == all_vars_no_dt
+
+
+def test_find_categorical_and_numerical_variables(df_vartypes):
+
+    # Case 1: user passes 1 variable that is categorical
+    assert find_categorical_and_numerical_variables(df_vartypes, ["Name"]) == (
+        ["Name"],
+        [],
+    )
+    assert find_categorical_and_numerical_variables(df_vartypes, "Name") == (
+        ["Name"],
+        [],
+    )
+
+    # Case 2: user passes 1 variable that is numerical
+    assert find_categorical_and_numerical_variables(df_vartypes, ["Age"]) == (
+        [],
+        ["Age"],
+    )
+    assert find_categorical_and_numerical_variables(df_vartypes, "Age") == (
+        [],
+        ["Age"],
+    )
+
+    # Case 3: user passes 1 categorical and 1 numerical variable
+    assert find_categorical_and_numerical_variables(df_vartypes, ["Age", "Name"]) == (
+        ["Name"],
+        ["Age"],
+    )
+
+    # Case 4: automatically identify variables
+    assert find_categorical_and_numerical_variables(df_vartypes, None) == (
+        ["Name", "City"],
+        ["Age", "Marks"],
+    )
+    assert find_categorical_and_numerical_variables(
+        df_vartypes[["Name", "City"]], None
+    ) == (["Name", "City"], [])
+    assert find_categorical_and_numerical_variables(
+        df_vartypes[["Age", "Marks"]], None
+    ) == ([], ["Age", "Marks"])
+
+    # Case 5: error when no variable is numerical or categorical
+    with pytest.raises(TypeError):
+        find_categorical_and_numerical_variables(df_vartypes["dob"].to_frame(), None)
+
+    with pytest.raises(TypeError):
+        find_categorical_and_numerical_variables(df_vartypes["dob"].to_frame(), ["dob"])
+
+    with pytest.raises(TypeError):
+        find_categorical_and_numerical_variables(df_vartypes["dob"].to_frame(), "dob")
+
+    # Case 6: user passes empty list
+    with pytest.raises(ValueError):
+        find_categorical_and_numerical_variables(df_vartypes, [])
+
+    # Case 7: datetime cast as object
+    df = df_vartypes.copy()
+    df["dob"] = df["dob"].astype("O")
+
+    # datetime variable is skipped when automatically finding variables, but
+    # selected if user passes it in list
+    assert find_categorical_and_numerical_variables(df, None) == (
+        ["Name", "City"],
+        ["Age", "Marks"],
+    )
+    assert find_categorical_and_numerical_variables(df, ["Name", "Marks", "dob"]) == (
+        ["Name", "dob"],
+        ["Marks"],
+    )
+
+    # Case 8: variables cast as category
+    df = df_vartypes.copy()
+    df["City"] = df["City"].astype("category")
+    assert find_categorical_and_numerical_variables(df, None) == (
+        ["Name", "City"],
+        ["Age", "Marks"],
+    )
+    assert find_categorical_and_numerical_variables(df, "City") == (["City"], [])