adds functionality to avoid selecting datetime variables as if they were categorical closes #336 (#344)

* categorical type checker no longer detects datetime

* works properly with numerical vars previously cast as object

* refactors typecheck condition

* rearranges tests, adds a few header comments

* removes datetime checks when variables is not None

* breaks down the fancy is_categorical logic, adds comments

* adds pedantic test, clarifies comment

* covers cases where dt var is cast as cat and vars are specified

Author: dodoarg

feature_engine/variable_manipulation.py

@@ -93,6 +93,25 @@ def _find_or_check_numerical_variables(
     return variables
 
 
+def _is_categorical_and_is_not_datetime(column: pd.Series) -> bool:
+
+    # check for datetime only if object cannot be cast as numeric because
+    # if it could pd.to_datetime would convert it to datetime regardless
+    if is_object(column):
+        is_categorical_and_is_not_datetime = is_numeric(
+            pd.to_numeric(column, errors="ignore")
+        ) or not is_datetime(pd.to_datetime(column, errors="ignore", utc=True))
+
+    # check for datetime only if the type of the categories is not numeric
+    # because pd.to_datetime throws an error when it is an integer
+    elif is_categorical(column):
+        is_categorical_and_is_not_datetime = is_numeric(
+            column.dtype.categories
+        ) or not is_datetime(pd.to_datetime(column, errors="ignore", utc=True))
+
+    return is_categorical_and_is_not_datetime
+
+
 def _find_or_check_categorical_variables(
     X: pd.DataFrame, variables: Variables = None
 ) -> List[Union[str, int]]:
@@ -119,7 +138,11 @@ def _find_or_check_categorical_variables(
 
     if variables is None:
         # find categorical variables in dataset
-        variables = list(X.select_dtypes(include=["O", "category"]).columns)
+        variables = [
+            column
+            for column in X.select_dtypes(include=["O", "category"]).columns
+            if _is_categorical_and_is_not_datetime(X[column])
+        ]
         if len(variables) == 0:
             raise ValueError(
                 "No categorical variables found in this dataframe. Please check "
@@ -136,7 +159,7 @@ def _find_or_check_categorical_variables(
         if len(variables) == 0:
             raise ValueError("The list of variables is empty.")
 
-        # check that user entered variables are of type numerical
+        # check that user entered variables are of type categorical
         else:
             if len(X[variables].select_dtypes(exclude=["O", "category"]).columns) > 0:
                 raise TypeError(

tests/test_variable_manipulation.py

@@ -68,42 +68,102 @@ def test_find_or_check_numerical_variables(df_vartypes, df_numeric_columns):
     assert _find_or_check_numerical_variables(df_numeric_columns, 2) == [2]
 
 
-def test_find_or_check_categorical_variables(df_vartypes, df_numeric_columns):
+def test_find_or_check_categorical_variables(
+    df_vartypes, df_datetime, df_numeric_columns
+):
     vars_cat = ["Name", "City"]
     vars_mix = ["Age", "Marks", "Name"]
 
-    assert _find_or_check_categorical_variables(df_vartypes, vars_cat) == vars_cat
-    assert _find_or_check_categorical_variables(df_vartypes, None) == vars_cat
-
+    # errors when vars entered by user are not categorical
     with pytest.raises(TypeError):
         assert _find_or_check_categorical_variables(df_vartypes, "Marks")
-
+    with pytest.raises(TypeError):
+        assert _find_or_check_categorical_variables(df_datetime, "datetime_range")
+    with pytest.raises(TypeError):
+        assert _find_or_check_categorical_variables(df_datetime, ["datetime_range"])
     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)
 
+    # error when user enters empty list
     with pytest.raises(ValueError):
         assert _find_or_check_categorical_variables(df_vartypes, variables=[])
 
+    # error when df has no categorical variables
     with pytest.raises(ValueError):
         assert _find_or_check_categorical_variables(df_vartypes[["Age", "Marks"]], None)
+    with pytest.raises(ValueError):
+        assert _find_or_check_categorical_variables(
+            df_datetime[["date_obj1", "time_obj"]], None
+        )
+
+    # when variables=None
+    assert _find_or_check_categorical_variables(df_vartypes, None) == vars_cat
+    assert _find_or_check_categorical_variables(df_datetime, None) == ["Name"]
+
+    # when vars are specified
+    assert _find_or_check_categorical_variables(df_vartypes, "Name") == ["Name"]
+    assert _find_or_check_categorical_variables(df_datetime, "date_obj1") == [
+        "date_obj1"
+    ]
+    assert _find_or_check_categorical_variables(df_vartypes, vars_cat) == vars_cat
+    assert _find_or_check_categorical_variables(df_datetime, ["Name", "date_obj1"]) == [
+        "Name",
+        "date_obj1",
+    ]
 
+    # vars specified, column name is integer
     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]
 
+    # numeric var cast as category
     df_vartypes["Age"] = df_vartypes["Age"].astype("category")
+    assert _find_or_check_categorical_variables(df_vartypes, "Age") == ["Age"]
     assert _find_or_check_categorical_variables(df_vartypes, None) == vars_cat + ["Age"]
     assert _find_or_check_categorical_variables(df_vartypes, ["Name", "Age"]) == [
         "Name",
         "Age",
     ]
 
+    # datetime vars cast as category
+    # object-like datetime
+    df_datetime["date_obj1"] = df_datetime["date_obj1"].astype("category")
+    assert _find_or_check_categorical_variables(df_datetime, None) == ["Name"]
+    assert _find_or_check_categorical_variables(df_datetime, ["Name", "date_obj1"]) == [
+        "Name",
+        "date_obj1",
+    ]
+    # datetime64
+    df_datetime["datetime_range"] = df_datetime["datetime_range"].astype("category")
+    assert _find_or_check_categorical_variables(df_datetime, None) == ["Name"]
+    assert _find_or_check_categorical_variables(
+        df_datetime, ["Name", "datetime_range"]
+    ) == ["Name", "datetime_range"]
+
+    # numeric var cast as object
+    df_vartypes["Marks"] = df_vartypes["Marks"].astype("O")
+    assert _find_or_check_categorical_variables(df_vartypes, "Marks") == ["Marks"]
+    assert _find_or_check_categorical_variables(df_vartypes, ["Name", "Marks"]) == [
+        "Name",
+        "Marks",
+    ]
+    assert _find_or_check_categorical_variables(df_vartypes, None) == vars_cat + [
+        "Age",
+        "Marks",
+    ]
+
+    # time-aware datetime var
+    tz_time = pd.DataFrame(
+        {"time_objTZ": df_datetime["time_obj"].add(["+5", "+11", "-3", "-8"])}
+    )
+    with pytest.raises(ValueError):
+        assert _find_or_check_categorical_variables(tz_time, None)
+    assert _find_or_check_categorical_variables(tz_time, "time_objTZ") == ["time_objTZ"]
+
 
 def test_find_or_check_datetime_variables(df_datetime):
     var_dt = ["datetime_range"]
@@ -116,7 +176,7 @@ def test_find_or_check_datetime_variables(df_datetime):
         {"date_obj1_cat": df_datetime["date_obj1"].astype("category")}
     )
     tz_time = pd.DataFrame(
-        {"time_objTZ": df_datetime["time_obj"].add(['+5', '+11', '-3', '-8'])}
+        {"time_objTZ": df_datetime["time_obj"].add(["+5", "+11", "-3", "-8"])}
     )
 
     # error when df has no datetime variables