Repository: feature-engine/feature_engine
Pull request #811 (2024-09-08 17:02:19)

add tz aware columns to variable handling functions
#810 

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/feature-engine/feature_engine/pull/811?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=feature-engine) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 97.99%. Comparing base [(`3e7b50e`)](https://app.codecov.io/gh/feature-engine/feature_engine/commit/3e7b50e908e6748ea3cd97ed0f414643552fdb51?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=feature-engine) to head [(`5b56324`)](https://app.codecov.io/gh/feature-engine/feature_engine/commit/5b56324ea7fbe7b35798b5fe2f68c3cf6ab6c370?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=feature-engine).

<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff           @@
##             main     #811   +/-   ##
=======================================
  Coverage   97.99%   97.99%           
=======================================
  Files         107      108    +1     
  Lines        4332     4335    +3     
  Branches      857      857           
=======================================
+ Hits         4245     4248    +3     
  Misses         54       54           
  Partials       33       33           
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/feature-engine/feature_engine/pull/811?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=feature-engine).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=feature-engine).


Patch:
diff --git a/feature_engine/variable_handling/check_variables.py b/feature_engine/variable_handling/check_variables.py
index 180d0d5..76c4ea7 100644
--- a/feature_engine/variable_handling/check_variables.py
+++ b/feature_engine/variable_handling/check_variables.py
@@ -8,6 +8,7 @@ from pandas.api.types import is_numeric_dtype as is_numeric
 from feature_engine.variable_handling._variable_type_checks import (
     _is_categorical_and_is_datetime,
 )
+from feature_engine.variable_handling.dtypes import DATETIME_TYPES
 
 Variables = Union[int, str, List[Union[str, int]]]
 
@@ -111,7 +112,8 @@ def check_datetime_variables(
     variables: Variables,
 ) -> List[Union[str, int]]:
     """
-    Checks that the variables in the list are or can be parsed as datetime.
+    Checks that the variables in the list are or can be parsed as datetime and or
+    datetimetz.
 
     More details in the :ref:`User Guide <check_datetime_vars>`.
 
@@ -147,7 +149,7 @@ def check_datetime_variables(
 
     # find non datetime variables, if any:
     non_datetime_vars = []
-    for column in X[variables].select_dtypes(exclude="datetime"):
+    for column in X[variables].select_dtypes(exclude=DATETIME_TYPES):
         if is_numeric(X[column]) or not _is_categorical_and_is_datetime(X[column]):
             non_datetime_vars.append(column)
 
diff --git a/feature_engine/variable_handling/dtypes.py b/feature_engine/variable_handling/dtypes.py
new file mode 100644
index 0000000..c7d9395
--- /dev/null
+++ b/feature_engine/variable_handling/dtypes.py
@@ -0,0 +1,1 @@
+DATETIME_TYPES = ("datetimetz", "datetime")
diff --git a/feature_engine/variable_handling/find_variables.py b/feature_engine/variable_handling/find_variables.py
index 6446bf0..04779ad 100644
--- a/feature_engine/variable_handling/find_variables.py
+++ b/feature_engine/variable_handling/find_variables.py
@@ -11,6 +11,7 @@ from feature_engine.variable_handling._variable_type_checks import (
     _is_categorical_and_is_datetime,
     _is_categorical_and_is_not_datetime,
 )
+from feature_engine.variable_handling.dtypes import DATETIME_TYPES
 
 
 def find_numerical_variables(X: pd.DataFrame) -> List[Union[str, int]]:
@@ -147,7 +148,7 @@ def find_all_variables(
 ) -> List[Union[str, int]]:
     """
     Returns a list with the names of all the variables in the dataframe. It has the
-    option to exlcude variables that can be parsed as datetime.
+    option to exlcude variables that can be parsed as datetime or datetimetz.
 
     More details in the :ref:`User Guide <find_all_vars>`.
 
@@ -178,7 +179,7 @@ def find_all_variables(
     ['var_num', 'var_cat', 'var_date']
     """
     if exclude_datetime is True:
-        variables = X.select_dtypes(exclude="datetime").columns.to_list()
+        variables = X.select_dtypes(exclude=DATETIME_TYPES).columns.to_list()
         variables = [
             var
             for var in variables
@@ -236,7 +237,6 @@ def find_categorical_and_numerical_variables(
 
     # 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 = []


Test patch:
diff --git a/tests/test_variable_handling/conftest.py b/tests/test_variable_handling/conftest.py
index ac4ad22..841656d 100644
--- a/tests/test_variable_handling/conftest.py
+++ b/tests/test_variable_handling/conftest.py
@@ -12,6 +12,9 @@ def df():
             "Marks": [0.9, 0.8, 0.7, 0.6],
             "date_range": pd.date_range("2020-02-24", periods=4, freq="min"),
             "date_obj0": ["2020-02-24", "2020-02-25", "2020-02-26", "2020-02-27"],
+            "date_range_tz": pd.date_range(
+                "2020-02-24", periods=4, freq="min"
+            ).tz_localize("UTC"),
         }
     )
     df["Name"] = df["Name"].astype("category")
@@ -21,7 +24,7 @@ def df():
 @pytest.fixture
 def df_int(df):
     df = df.copy()
-    df.columns = [1, 2, 3, 4, 5, 6]
+    df.columns = range(1, len(df.columns) + 1)
     return df
 
 
diff --git a/tests/test_variable_handling/test_check_variables.py b/tests/test_variable_handling/test_check_variables.py
index 8202ff7..8eba88c 100644
--- a/tests/test_variable_handling/test_check_variables.py
+++ b/tests/test_variable_handling/test_check_variables.py
@@ -99,6 +99,7 @@ def test_check_datetime_variables_returns_datetime_variables(df_datetime):
     vars_convertible_to_dt = ["date_range", "date_obj1", "date_obj2", "time_obj"]
     var_convertible_to_dt = "date_obj1"
     tz_time = "time_objTZ"
+    tz_time_obj = "date_range_tz"
 
     # when variables are specified
     assert check_datetime_variables(df_datetime, var_dt_str) == [var_dt_str]
@@ -111,6 +112,7 @@ def test_check_datetime_variables_returns_datetime_variables(df_datetime):
         == vars_convertible_to_dt
     )
     assert check_datetime_variables(df_datetime, tz_time) == [tz_time]
+    assert check_datetime_variables(df_datetime, tz_time_obj) == [tz_time_obj]
 
     df_datetime[vars_convertible_to_dt] = df_datetime[vars_convertible_to_dt].astype(
         pd.CategoricalDtype
diff --git a/tests/test_variable_handling/test_find_variables.py b/tests/test_variable_handling/test_find_variables.py
index 06b53d6..efb9b53 100644
--- a/tests/test_variable_handling/test_find_variables.py
+++ b/tests/test_variable_handling/test_find_variables.py
@@ -51,6 +51,7 @@ def test_datetime_variables_finds_datetime_variables(df_datetime):
     vars_dt = [
         "date_range",
         "date_obj0",
+        "date_range_tz",
         "date_obj1",
         "date_obj2",
         "time_obj",
@@ -65,7 +66,6 @@ def test_datetime_variables_finds_datetime_variables(df_datetime):
 
 
 def test_datetime_variables_raises_error_when_no_datetime_variables(df_datetime):
-
     msg = "No datetime variables found in this dataframe."
 
     vars_nondt = ["Marks", "Age", "Name"]
@@ -76,7 +76,15 @@ def test_datetime_variables_raises_error_when_no_datetime_variables(df_datetime)
 
 
 def test_find_all_variables(df):
-    all_vars = ["Name", "City", "Age", "Marks", "date_range", "date_obj0"]
+    all_vars = [
+        "Name",
+        "City",
+        "Age",
+        "Marks",
+        "date_range",
+        "date_obj0",
+        "date_range_tz",
+    ]
     all_vars_no_dt = ["Name", "City", "Age", "Marks"]
 
     assert find_all_variables(df, exclude_datetime=False) == all_vars
@@ -84,7 +92,6 @@ def test_find_all_variables(df):
 
 
 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"],