finalise migration of variable handling module

Author: solegalli

AGENTS.md

@@ -0,0 +1,55 @@
+# AGENTS.md
+
+Conventions for working in this repo. Optimize for readability and speed, 
+in that order of how you decide, but don't ship a slow default when a 
+fast one is free.
+
+## Inputs
+
+Feature-engine transformers take dataframes (pandas, polars, or any other
+narwhals-supported backend) as input, not numpy arrays. Don't add
+handling for array input.
+
+## Booleans and control flow
+
+- Compare booleans explicitly: `if x is True:` / `if x is False:`, never
+  `if x:` / `if not x:`.
+- Check container emptiness with `len(x) == 0`, never `if not x:`.
+- `isinstance(...)` checks and `in`/`not in` membership tests are already
+  explicit — leave them as-is, this rule isn't about those.
+
+## Comments
+
+Max 2 lines. Only explain a non-obvious WHY (a hidden constraint, a subtle
+backend difference, a workaround) — never describe WHAT the code does.
+
+## Don't anticipate errors
+
+Don't add error handling or validation for scenarios that can't happen. If
+unsure whether something can happen, check it (grep, run a quick repro) or
+ask — don't guess and defensively code around it.
+
+## Redundant lists/sets
+
+- Narwhals' `.columns` is already `list[str]` — don't wrap it in `list()`.
+- pandas' `.columns` is an `Index`, not a list — `list()` is required there
+  (an `Index == list` comparison is elementwise, not a clean bool).
+
+## Verify before applying
+
+Benchmark before claiming a speedup, and diff old-vs-new output across
+realistic and edge cases (empty/all-NaN, both backends, both dtype
+branches) before trusting a rewrite — logic mistakes here are easy to make
+and easy to miss without an actual comparison.
+
+## Tests
+
+- `pytest.raises(ExceptionType, match=msg)`, never
+  `with pytest.raises() as record: ... assert str(record.value) == msg`.
+
+## API changes
+
+- New parameters default to preserve current behavior.
+- When adding a parameter to a function called from multiple sites (or a
+  shared private helper), thread it through every call site, not just the
+  one you're looking at.

feature_engine/variable_handling/_variable_type_checks.py

@@ -1,3 +1,4 @@
+import warnings
 from datetime import date, datetime
 
 import narwhals as nw
@@ -31,18 +32,35 @@ def _is_convertible_to_num(s: "nw.Series") -> bool:
     if len(values) == 0:
         return False
     try:
-        for value in values:
+        for value in values[:100]:
             float(value)
     except (ValueError, TypeError):
         return False
     return True
 
 
 def _is_convertible_to_dt(s: "nw.Series") -> bool:
-    values = s.drop_nulls().to_list()
-    if len(values) == 0:
+    values = s.drop_nulls()
+    values_list = values.to_list()
+    if len(values_list) == 0:
         return False
-    for value in values:
+
+    first_value = values_list[0]
+    if not isinstance(first_value, (date, datetime)):
+        if _looks_like_date_string(first_value) is False:
+            return False
+
+    # Try the backend's own vectorized parser first (faster).
+    # Fall back to the per-value check (below) when it fails.
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", UserWarning)
+        try:
+            values.str.to_datetime()
+            return True
+        except Exception:
+            pass
+
+    for value in values_list[:100]:
         if isinstance(value, (date, datetime)):
             continue
         if _looks_like_date_string(value) is False:

feature_engine/variable_handling/dtypes.py

@@ -34,7 +34,7 @@ def _find_nw_categoricals(
             | nw.selectors.string()
             | nw.selectors.by_dtype(nw.Object)
         )
-        # `|`-combined selectors don't preserve the dataframe's column order,
+        # `|`-combined selectors don't preserve column order,
         # so re-filter over nw_X.columns to restore it.
         matched = set(nw_X.select(_NW_SELECTOR).columns)
         candidates = [column for column in nw_X.columns if column in matched]

feature_engine/variable_handling/find_variables.py

@@ -3,6 +3,7 @@
 from typing import List, Union
 
 import narwhals as nw
+import narwhals.dependencies as nwd
 from narwhals.typing import IntoDataFrame
 
 Variables = Union[int, str, List[Union[str, int]]]
@@ -43,7 +44,10 @@ def retain_variables_if_in_df(X: IntoDataFrame, variables):
     if isinstance(variables, (str, int)):
         variables = [variables]
 
-    columns = nw.from_native(X, eager_only=True).columns
+    if nwd.is_pandas_dataframe(X) is True:
+        columns = set(X.columns)
+    else:
+        columns = set(nw.from_native(X, eager_only=True).columns)
     variables_in_df = [var for var in variables if var in columns]
 
     # Raise an error if no column is left to work with.

feature_engine/variable_handling/retain_variables.py

@@ -17,22 +17,14 @@ def cast_categorical(df, columns):
     return df.with_columns([pl.col(c).cast(pl.Categorical) for c in columns])
 
 
-# Data shared between the pandas and polars variants of a test. Kept as plain
-# dicts/lists (not fixtures) so a test can build both `make_df(BASIC_DATA)` and
-# `make_df(BASIC_DATA)` for a different backend without needing to convert
-# between frame types.
+# Data shared between the pandas and polars variants of a test.
 BASIC_DATA = {
     "Name": ["tom", "nick", "krish", "jack"],
     "City": ["London", "Manchester", "Liverpool", "Bristol"],
     "Age": [20, 21, 19, 18],
     "Marks": [0.9, 0.8, 0.7, 0.6],
 }
 
-# Datetime formats that both pandas and polars auto-detect: native
-# Datetime/Date columns and ISO-8601 strings. Formats that only pandas'
-# flexible, dateutil-backed guessing can parse (e.g. "01-Jan-2010",
-# "10/11/12", bare time strings) are exercised separately, in pandas-only
-# tests, against the `df_datetime` fixture below.
 DATETIME_DATA = {
     **BASIC_DATA,
     "date_range": [datetime(2020, 2, 24, 0, i) for i in range(4)],

tests/test_variable_handling/conftest.py

@@ -129,7 +129,7 @@ def test_check_datetime_variables_returns_datetime_variables(make_df):
     assert check_datetime_variables(df, vars_dt) == vars_dt
     assert check_datetime_variables(df, tz_time) == [tz_time]
 
-    # only the string column can be cast to categorical - native Datetime
+    # only the string column can be cast to categorical. Native Datetime
     # columns can't be cast to Categorical in polars
     df = cast_categorical(df, ["date_obj0"])
     assert check_datetime_variables(df, "date_obj0") == ["date_obj0"]

tests/test_variable_handling/test_check_variables.py

@@ -25,7 +25,7 @@ def test_numerical_variables_finds_variables(make_df):
 
 
 def test_numerical_variables_finds_variables_with_int_column_names(df_int):
-    # polars requires string column names, so int-named columns are pandas-only
+    # polars requires string column names. int-named columns are pandas-only
     assert find_numerical_variables(df_int) == [3, 4]
 
 
@@ -336,11 +336,8 @@ def test_numcat_vars_as_category(make_df):
 
 @pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame])
 def test_numcat_agrees_with_find_categorical_on_date_like_category(make_df):
-    # Regression test: the single-variable branch of
-    # find_categorical_and_numerical_variables used to short-circuit past the
-    # datetime check for Categorical-dtype columns, so it disagreed with
-    # find_categorical_variables on a category column holding date-like
-    # strings. All three calls below must now exclude it consistently.
+    # Regression test: the single-variable path used to disagree with
+    # find_categorical_variables on a date-like category column.
     df = make_df({"date_cat": DATETIME_DATA["date_obj0"], "num": BASIC_DATA["Age"]})
     df = cast_categorical(df, ["date_cat"])
 

tests/test_variable_handling/test_find_variables.py

@@ -4,46 +4,36 @@
 
 from feature_engine.variable_handling.retain_variables import retain_variables_if_in_df
 
-
-def make_empty_df(is_pandas: bool, columns):
-    if is_pandas:
-        return pd.DataFrame(columns=columns)
-    return pl.DataFrame(schema=columns)
-
-
 test_dict = [
     (["A", "C", "B", "G", "H"], ["A", "C", "B"], ["X", "Y"]),
     ("C", ["C"], "G"),
 ]
 
 
-@pytest.mark.parametrize("is_pandas", [True, False])
+@pytest.mark.parametrize("make_df", [pd.DataFrame, pl.DataFrame])
 @pytest.mark.parametrize("variables, overlap, col_not_in_df", test_dict)
-def test_retain_variables_if_in_df(is_pandas, variables, overlap, col_not_in_df):
-    df = make_empty_df(is_pandas, ["A", "B", "C", "D", "E"])
+def test_retain_variables_if_in_df(make_df, variables, overlap, col_not_in_df):
+    df = make_df({"A": [1], "B": [1], "C": [1], "D": [1], "E": [1]})
 
     msg = "None of the variables in the list are present in the dataframe."
 
     assert retain_variables_if_in_df(df, variables) == overlap
 
-    with pytest.raises(ValueError) as record:
+    with pytest.raises(ValueError, match=msg):
         retain_variables_if_in_df(df, col_not_in_df)
-    assert str(record.value) == msg
 
 
 def test_retain_variables_if_in_df_int_column_names():
-    # polars requires string column names, so int-named columns are pandas-only
-    df = pd.DataFrame(columns=[1, 2, 3, 4, 5])
+    # polars requires string column names. int-named columns are pandas-only
+    df = pd.DataFrame({1: [1], 2: [1], 3: [1], 4: [1], 5: [1]})
 
     msg = "None of the variables in the list are present in the dataframe."
 
     assert retain_variables_if_in_df(df, [1, 2, 4, 6]) == [1, 2, 4]
     assert retain_variables_if_in_df(df, 1) == [1]
 
-    with pytest.raises(ValueError) as record:
+    with pytest.raises(ValueError, match=msg):
         retain_variables_if_in_df(df, [6, 7])
-    assert str(record.value) == msg
 
-    with pytest.raises(ValueError) as record:
+    with pytest.raises(ValueError, match=msg):
         retain_variables_if_in_df(df, 7)
-    assert str(record.value) == msg