Repository: feature-engine/feature_engine

Issue:
categorical encoders should not pick datetime variables automatically
If a datetime variable is encoded as object or categorical, and the parameter `variables` is set to None, the transformer will detect those datetime variables as categorical and then encode them.

It might be better practice to exclude datetime variables when the transformer searches for variables automatically


Discussion:
dodoarg: I'm glad you appreciate my contribution, I feel like I have learned a lot considering where I started so thank you for your valuable feedback

Yes I agree on giving the user the last word on what variables should be considered categorical if they specify them. Not sure the way I've broken down the logic is exactly what you meant but it should look better now. What do you think?


> just go ahead and add code to tackle the 2 problematic cases for the datetime variable checker function as well

Just to make sure, you mean in this PR?

Patch:
diff --git a/feature_engine/outliers/base_outlier.py b/feature_engine/outliers/base_outlier.py
index 881d201..a57caf2 100644
--- a/feature_engine/outliers/base_outlier.py
+++ b/feature_engine/outliers/base_outlier.py
@@ -2,6 +2,7 @@ import numpy as np
 import pandas as pd
 from sklearn.base import BaseEstimator, TransformerMixin
 from sklearn.utils.validation import check_is_fitted
+from typing import List, Optional, Union
 
 from feature_engine.dataframe_checks import (
     _check_contains_inf,
@@ -10,6 +11,10 @@ from feature_engine.dataframe_checks import (
     _is_dataframe,
 )
 from feature_engine.validation import _return_tags
+from feature_engine.variable_manipulation import (
+    _check_input_parameter_variables,
+    _find_or_check_numerical_variables,
+)
 
 
 class BaseOutlier(BaseEstimator, TransformerMixin):
@@ -89,3 +94,108 @@ class BaseOutlier(BaseEstimator, TransformerMixin):
 
     def _more_tags(self):
         return _return_tags()
+
+
+class WinsorizerBase(BaseOutlier):
+    def __init__(
+        self,
+        capping_method: str = "gaussian",
+        tail: str = "right",
+        fold: Union[int, float] = 3,
+        variables: Union[None, int, str, List[Union[str, int]]] = None,
+        missing_values: str = "raise",
+    ) -> None:
+
+        if capping_method not in ["gaussian", "iqr", "quantiles"]:
+            raise ValueError(
+                "capping_method takes only values 'gaussian', 'iqr' or 'quantiles'"
+            )
+
+        if tail not in ["right", "left", "both"]:
+            raise ValueError("tail takes only values 'right', 'left' or 'both'")
+
+        if fold <= 0:
+            raise ValueError("fold takes only positive numbers")
+
+        if capping_method == "quantiles" and fold > 0.2:
+            raise ValueError(
+                "with capping_method ='quantiles', fold takes values between 0 and "
+                "0.20 only."
+            )
+
+        if missing_values not in ["raise", "ignore"]:
+            raise ValueError("missing_values takes only values 'raise' or 'ignore'")
+
+        self.capping_method = capping_method
+        self.tail = tail
+        self.fold = fold
+        self.variables = _check_input_parameter_variables(variables)
+        self.missing_values = missing_values
+
+    def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None):
+        """
+        Learn the values that should be used to replace outliers.
+
+        Parameters
+        ----------
+        X : pandas dataframe of shape = [n_samples, n_features]
+            The training input samples.
+
+        y : pandas Series, default=None
+            y is not needed in this transformer. You can pass y or None.
+        """
+
+        # check input dataframe
+        X = _is_dataframe(X)
+
+        # find or check for numerical variables
+        self.variables_ = _find_or_check_numerical_variables(X, self.variables)
+
+        if self.missing_values == "raise":
+            # check if dataset contains na
+            _check_contains_na(X, self.variables_)
+            _check_contains_inf(X, self.variables_)
+
+        self.right_tail_caps_ = {}
+        self.left_tail_caps_ = {}
+
+        # estimate the end values
+        if self.tail in ["right", "both"]:
+            if self.capping_method == "gaussian":
+                self.right_tail_caps_ = (
+                    X[self.variables_].mean() + self.fold * X[self.variables_].std()
+                ).to_dict()
+
+            elif self.capping_method == "iqr":
+                IQR = X[self.variables_].quantile(0.75) - X[self.variables_].quantile(
+                    0.25
+                )
+                self.right_tail_caps_ = (
+                    X[self.variables_].quantile(0.75) + (IQR * self.fold)
+                ).to_dict()
+
+            elif self.capping_method == "quantiles":
+                self.right_tail_caps_ = (
+                    X[self.variables_].quantile(1 - self.fold).to_dict()
+                )
+
+        if self.tail in ["left", "both"]:
+            if self.capping_method == "gaussian":
+                self.left_tail_caps_ = (
+                    X[self.variables_].mean() - self.fold * X[self.variables_].std()
+                ).to_dict()
+
+            elif self.capping_method == "iqr":
+                IQR = X[self.variables_].quantile(0.75) - X[self.variables_].quantile(
+                    0.25
+                )
+                self.left_tail_caps_ = (
+                    X[self.variables_].quantile(0.25) - (IQR * self.fold)
+                ).to_dict()
+
+            elif self.capping_method == "quantiles":
+                self.left_tail_caps_ = X[self.variables_].quantile(self.fold).to_dict()
+
+        self.n_features_in_ = X.shape[1]
+
+        return self
diff --git a/feature_engine/outliers/trimmer.py b/feature_engine/outliers/trimmer.py
index 0719241..6936586 100644
--- a/feature_engine/outliers/trimmer.py
+++ b/feature_engine/outliers/trimmer.py
@@ -4,10 +4,10 @@
 import numpy as np
 import pandas as pd
 
-from feature_engine.outliers import Winsorizer
+from feature_engine.outliers.base_outlier import WinsorizerBase
 
 
-class OutlierTrimmer(Winsorizer):
+class OutlierTrimmer(WinsorizerBase):
     """The OutlierTrimmer() removes observations with outliers from the dataset.
 
     The OutlierTrimmer() first calculates the maximum and /or minimum values
diff --git a/feature_engine/outliers/winsorizer.py b/feature_engine/outliers/winsorizer.py
index 23a620f..4d76aa9 100644
--- a/feature_engine/outliers/winsorizer.py
+++ b/feature_engine/outliers/winsorizer.py
@@ -1,26 +1,19 @@
 # Authors: Soledad Galli <solegalli@protonmail.com>
 # License: BSD 3 clause
 
-from typing import List, Optional, Union
+from typing import List, Union
 
+import numpy as np
 import pandas as pd
 
-from feature_engine.dataframe_checks import (
-    _check_contains_inf,
-    _check_contains_na,
-    _is_dataframe,
-)
-from feature_engine.outliers.base_outlier import BaseOutlier
-from feature_engine.variable_manipulation import (
-    _check_input_parameter_variables,
-    _find_or_check_numerical_variables,
-)
+from feature_engine.dataframe_checks import _is_dataframe
+from feature_engine.outliers.base_outlier import WinsorizerBase
 
 
-class Winsorizer(BaseOutlier):
+class Winsorizer(WinsorizerBase):
     """
     The Winsorizer() caps maximum and/or minimum values of a variable at automatically
-    determined values.
+    determined values, and optionally adds indicators.
 
     The values to cap variables are determined using:
 
@@ -95,6 +88,11 @@ class Winsorizer(BaseOutlier):
         both sides. Thus, when `capping_method='quantiles'`, then `'fold'` takes values
         between 0 and 0.20.
 
+    add_indicators: bool, default=False
+        Whether to add indicator variables to flag the capped outliers.
+        If 'True', binary variables will be added to flag outliers on the left and right
+        tails of the distribution. One binary variable per tail, per variable.
+
     variables: list, default=None
         The list of variables for which the outliers will be capped. If None,
         the transformer will select and cap all numerical variables.
@@ -136,108 +134,59 @@ class Winsorizer(BaseOutlier):
         capping_method: str = "gaussian",
         tail: str = "right",
         fold: Union[int, float] = 3,
+        add_indicators: bool = False,
         variables: Union[None, int, str, List[Union[str, int]]] = None,
         missing_values: str = "raise",
     ) -> None:
-
-        if capping_method not in ["gaussian", "iqr", "quantiles"]:
-            raise ValueError(
-                "capping_method takes only values 'gaussian', 'iqr' or 'quantiles'"
-            )
-
-        if tail not in ["right", "left", "both"]:
-            raise ValueError("tail takes only values 'right', 'left' or 'both'")
-
-        if fold <= 0:
-            raise ValueError("fold takes only positive numbers")
-
-        if capping_method == "quantiles" and fold > 0.2:
+        if not isinstance(add_indicators, bool):
             raise ValueError(
-                "with capping_method ='quantiles', fold takes values between 0 and "
-                "0.20 only."
+                "add_indicators takes only booleans True and False"
+                f"Got {add_indicators} instead."
             )
+        super().__init__(capping_method, tail, fold, variables, missing_values)
+        self.add_indicators = add_indicators
 
-        if missing_values not in ["raise", "ignore"]:
-            raise ValueError("missing_values takes only values 'raise' or 'ignore'")
-
-        self.capping_method = capping_method
-        self.tail = tail
-        self.fold = fold
-        self.variables = _check_input_parameter_variables(variables)
-        self.missing_values = missing_values
-
-    def fit(self, X: pd.DataFrame, y: Optional[pd.Series] = None):
+    def transform(self, X: pd.DataFrame) -> pd.DataFrame:
         """
-        Learn the values that should be used to replace outliers.
+        Cap the variable values. Optionally, add outlier indicators.
 
         Parameters
         ----------
-        X : pandas dataframe of shape = [n_samples, n_features]
-            The training input samples.
-
-        y : pandas Series, default=None
-            y is not needed in this transformer. You can pass y or None.
+        X: pandas dataframe of shape = [n_samples, n_features]
+            The data to be transformed.
+
+        Returns
+        -------
+        X_new: pandas dataframe of shape = [n_samples, n_features + n_ind]
+            The dataframe with the capped variables and indicators.
+            The number of output variables depends on the values for 'tail' and
+            'add_indicators': if passing 'add_indicators=False', will be equal
+            to 'n_features', otherwise, will have an additional indicator column
+            per processed feature for each tail.
         """
-
-        # check input dataframe
-        X = _is_dataframe(X)
-
-        # find or check for numerical variables
-        self.variables_ = _find_or_check_numerical_variables(X, self.variables)
-
-        if self.missing_values == "raise":
-            # check if dataset contains na
-            _check_contains_na(X, self.variables_)
-            _check_contains_inf(X, self.variables_)
-
-        self.right_tail_caps_ = {}
-        self.left_tail_caps_ = {}
-
-        # estimate the end values
-        if self.tail in ["right", "both"]:
-            if self.capping_method == "gaussian":
-                self.right_tail_caps_ = (
-                    X[self.variables_].mean() + self.fold * X[self.variables_].std()
-                ).to_dict()
-
-            elif self.capping_method == "iqr":
-                IQR = X[self.variables_].quantile(0.75) - X[self.variables_].quantile(
-                    0.25
-                )
-                self.right_tail_caps_ = (
-                    X[self.variables_].quantile(0.75) + (IQR * self.fold)
-                ).to_dict()
-
-            elif self.capping_method == "quantiles":
-                self.right_tail_caps_ = (
-                    X[self.variables_].quantile(1 - self.fold).to_dict()
-                )
-
-        if self.tail in ["left", "both"]:
-            if self.capping_method == "gaussian":
-                self.left_tail_caps_ = (
-                    X[self.variables_].mean() - self.fold * X[self.variables_].std()
-                ).to_dict()
-
-            elif self.capping_method == "iqr":
-                IQR = X[self.variables_].quantile(0.75) - X[self.variables_].quantile(
-                    0.25
-                )
-                self.left_tail_caps_ = (
-                    X[self.variables_].quantile(0.25) - (IQR * self.fold)
-                ).to_dict()
-
-            elif self.capping_method == "quantiles":
-                self.left_tail_caps_ = X[self.variables_].quantile(self.fold).to_dict()
-
-        self.n_features_in_ = X.shape[1]
-
-        return self
-
-    # Ugly work around to import the docstring for Sphinx, otherwise not necessary
-    def transform(self, X: pd.DataFrame) -> pd.DataFrame:
-        X = super().transform(X)
-
-        return X
-
-    transform.__doc__ = BaseOutlier.transform.__doc__
+        if not self.add_indicators:
+            return super().transform(X)
+        else:
+            X_orig = _is_dataframe(X)
+            X_out = super().transform(X_orig)
+            X_orig = X_orig[self.variables_]
+            X_out_filtered = X_out[self.variables_]
+
+            if self.tail in ["left", "both"]:
+                X_left = X_out_filtered > X_orig
+                X_left.columns = [str(cl) + "_left" for cl in self.variables_]
+            if self.tail in ["right", "both"]:
+                X_right = X_out_filtered < X_orig
+                X_right.columns = [str(cl) + "_right" for cl in self.variables_]
+            if self.tail == "left":
+                X_out = pd.concat([X_out, X_left.astype(np.float64)], axis=1)
+            elif self.tail == "right":
+                X_out = pd.concat([X_out, X_right.astype(np.float64)], axis=1)
+            else:
+                X_both = pd.concat([X_left, X_right], axis=1).astype(np.float64)
+                X_both = X_both[[
+                    cl1 for cl2 in zip(X_left.columns.values, X_right.columns.values)
+                    for cl1 in cl2
+                ]]
+                X_out = pd.concat([X_out, X_both], axis=1)
+        return X_out
diff --git a/feature_engine/variable_manipulation.py b/feature_engine/variable_manipulation.py
index bf8135c..4355e56 100644
--- a/feature_engine/variable_manipulation.py
+++ b/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(


Test patch:
diff --git a/tests/test_outliers/test_winsorizer.py b/tests/test_outliers/test_winsorizer.py
index 2b716a0..1737d9f 100644
--- a/tests/test_outliers/test_winsorizer.py
+++ b/tests/test_outliers/test_winsorizer.py
@@ -164,6 +164,87 @@ def test_quantile_capping_both_tails_with_fold_15_percent(df_normal_dist):
     assert np.round(df_normal_dist["var"].max(), 3) > np.round(0.11823196128033647, 3)
 
 
+def test_indicators_are_added(df_normal_dist):
+    transformer = Winsorizer(
+        tail="both", capping_method="quantiles", fold=0.1, add_indicators=True)
+    X = transformer.fit_transform(df_normal_dist)
+    # test that the number of output variables is correct
+    assert X.shape[1] == 3 * df_normal_dist.shape[1]
+    assert np.all(X.iloc[:, df_normal_dist.shape[1]:].sum(axis=0) > 0)
+
+    transformer = Winsorizer(
+        tail="left", capping_method="quantiles", fold=0.1, add_indicators=True)
+    X = transformer.fit_transform(df_normal_dist)
+    assert X.shape[1] == 2 * df_normal_dist.shape[1]
+    assert np.all(X.iloc[:, df_normal_dist.shape[1]:].sum(axis=0) > 0)
+
+    transformer = Winsorizer(
+        tail="right", capping_method="quantiles", fold=0.1, add_indicators=True)
+    X = transformer.fit_transform(df_normal_dist)
+    assert X.shape[1] == 2 * df_normal_dist.shape[1]
+    assert np.all(X.iloc[:, df_normal_dist.shape[1]:].sum(axis=0) > 0)
+
+
+def test_indicators_filter_variables(df_vartypes):
+    transformer = Winsorizer(
+        variables=["Age", "Marks"],
+        tail="both",
+        capping_method="quantiles",
+        fold=0.1,
+        add_indicators=True
+    )
+    X = transformer.fit_transform(df_vartypes)
+    assert X.shape[1] == df_vartypes.shape[1] + 4
+
+    transformer.set_params(tail="left")
+    X = transformer.fit_transform(df_vartypes)
+    assert X.shape[1] == df_vartypes.shape[1] + 2
+
+    transformer.set_params(tail="right")
+    X = transformer.fit_transform(df_vartypes)
+    assert X.shape[1] == df_vartypes.shape[1] + 2
+
+
+def test_indicators_are_correct():
+    transformer = Winsorizer(
+        tail="left",
+        capping_method="quantiles",
+        fold=0.1,
+        add_indicators=True
+    )
+    df = pd.DataFrame({"col": np.arange(100).astype(np.float64)})
+    df_out = transformer.fit_transform(df)
+    expected_ind = np.r_[np.repeat(True, 10), np.repeat(False, 90)].astype(np.float64)
+    pd.testing.assert_frame_equal(
+        df_out.drop("col", axis=1),
+        df.assign(col_left=expected_ind).drop("col", axis=1)
+    )
+
+    transformer.set_params(tail="right")
+    df_out = transformer.fit_transform(df)
+    expected_ind = np.r_[np.repeat(False, 90), np.repeat(True, 10)].astype(np.float64)
+    pd.testing.assert_frame_equal(
+        df_out.drop("col", axis=1),
+        df.assign(col_right=expected_ind).drop("col", axis=1)
+    )
+
+    transformer.set_params(tail="both")
+    df_out = transformer.fit_transform(df)
+    expected_ind_left = np.r_[
+        np.repeat(True, 10), np.repeat(False, 90)
+    ].astype(np.float64)
+    expected_ind_right = np.r_[
+        np.repeat(False, 90), np.repeat(True, 10)
+    ].astype(np.float64)
+    pd.testing.assert_frame_equal(
+        df_out.drop("col", axis=1),
+        df.assign(
+            col_left=expected_ind_left,
+            col_right=expected_ind_right
+        ).drop("col", axis=1)
+    )
+
+
 def test_transformer_ignores_na_in_df(df_na):
     # test case 7: dataset contains na and transformer is asked to ignore them
     transformer = Winsorizer(
@@ -227,6 +308,15 @@ def test_error_if_capping_method_quantiles_and_fold_value_not_permitted():
         Winsorizer(capping_method="quantiles", fold=0.3)
 
 
+def test_error_if_add_incators_not_permitted():
+    with pytest.raises(ValueError):
+        Winsorizer(add_indicators=-1)
+    with pytest.raises(ValueError):
+        Winsorizer(add_indicators=())
+    with pytest.raises(ValueError):
+        Winsorizer(add_indicators=[True])
+
+
 def test_fit_raises_error_if_na_in_inut_df(df_na):
     # test case 8: when dataset contains na, fit method
     with pytest.raises(ValueError):
diff --git a/tests/test_variable_manipulation.py b/tests/test_variable_manipulation.py
index 635a9bb..bd0302f 100644
--- a/tests/test_variable_manipulation.py
+++ b/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