fix: nw.selectors.by_dtype(nw.Datetime) was excluding tz-aware for Polars

narwhals/_polars/namespace.py

@@ -13,6 +13,7 @@
 from narwhals._expression_parsing import parse_into_exprs
 from narwhals._polars.utils import extract_args_kwargs
 from narwhals._polars.utils import narwhals_to_native_dtype
+from narwhals.dtypes import DType
 from narwhals.utils import Implementation
 
 if TYPE_CHECKING:
@@ -22,7 +23,6 @@
     from narwhals._polars.dataframe import PolarsLazyFrame
     from narwhals._polars.expr import PolarsExpr
     from narwhals._polars.typing import IntoPolarsExpr
-    from narwhals.dtypes import DType
     from narwhals.utils import Version
 
 
@@ -219,10 +219,15 @@ def by_dtype(self: Self, dtypes: Iterable[DType]) -> PolarsExpr:
 
         from narwhals._polars.expr import PolarsExpr
 
+        native_dtypes: list[pl.DataType | type[pl.DataType]] = []
+        for dtype in dtypes:
+            native_dtype_instantiated = narwhals_to_native_dtype(dtype, self._version)
+            if isinstance(dtype, type) and issubclass(dtype, DType):
+                native_dtypes.append(native_dtype_instantiated.__class__)
+            else:
+                native_dtypes.append(native_dtype_instantiated)
         return PolarsExpr(
-            pl.selectors.by_dtype(
-                [narwhals_to_native_dtype(dtype, self._version) for dtype in dtypes]
-            ),
+            pl.selectors.by_dtype(native_dtypes),
             version=self._version,
             backend_version=self._backend_version,
         )

tests/selectors_test.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import re
+from datetime import datetime
 
 import pytest
 
@@ -125,3 +126,17 @@ def test_set_ops_invalid(constructor: Constructor) -> None:
         match=re.escape("unsupported operand type(s) for op: ('Selector' + 'Selector')"),
     ):
         df.select(boolean() + numeric())
+
+
+def test_tz_aware(constructor: Constructor, request: pytest.FixtureRequest) -> None:
+    if "polars" in str(constructor) and POLARS_VERSION < (1, 19):
+        # bug in old polars
+        request.applymarker(pytest.mark.xfail)
+
+    data = {"a": [datetime(2020, 1, 1), datetime(2020, 1, 2)], "c": [4, 5]}
+    df = nw.from_native(constructor(data)).with_columns(
+        b=nw.col("a").dt.replace_time_zone("Asia/Katmandu")
+    )
+    result = df.select(nw.selectors.by_dtype(nw.Datetime)).collect_schema().names()
+    expected = ["a", "b"]
+    assert result == expected