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
 
 
@@ -35,8 +35,6 @@ def __init__(
         self._version = version
 
     def __getattr__(self: Self, attr: str) -> Any:
-        import polars as pl
-
         from narwhals._polars.expr import PolarsExpr
 
         def func(*args: Any, **kwargs: Any) -> Any:
@@ -50,8 +48,6 @@ def func(*args: Any, **kwargs: Any) -> Any:
         return func
 
     def nth(self: Self, *indices: int) -> PolarsExpr:
-        import polars as pl
-
         from narwhals._polars.expr import PolarsExpr
 
         if self._backend_version < (1, 0, 0):
@@ -62,8 +58,6 @@ def nth(self: Self, *indices: int) -> PolarsExpr:
         )
 
     def len(self: Self) -> PolarsExpr:
-        import polars as pl
-
         from narwhals._polars.expr import PolarsExpr
 
         if self._backend_version < (0, 20, 5):
@@ -114,8 +108,6 @@ def concat(
         )
 
     def lit(self: Self, value: Any, dtype: DType | None) -> PolarsExpr:
-        import polars as pl
-
         from narwhals._polars.expr import PolarsExpr
 
         if dtype is not None:
@@ -129,8 +121,6 @@ def lit(self: Self, value: Any, dtype: DType | None) -> PolarsExpr:
         )
 
     def mean_horizontal(self: Self, *exprs: IntoPolarsExpr) -> PolarsExpr:
-        import polars as pl
-
         from narwhals._polars.expr import PolarsExpr
 
         polars_exprs = cast("list[PolarsExpr]", parse_into_exprs(*exprs, namespace=self))
@@ -155,8 +145,6 @@ def concat_str(
         separator: str,
         ignore_nulls: bool,
     ) -> PolarsExpr:
-        import polars as pl
-
         from narwhals._polars.expr import PolarsExpr
 
         pl_exprs: list[pl.Expr] = [
@@ -215,14 +203,16 @@ def __init__(self: Self, version: Version, backend_version: tuple[int, ...]) ->
         self._backend_version = backend_version
 
     def by_dtype(self: Self, dtypes: Iterable[DType]) -> PolarsExpr:
-        import polars as pl
-
         from narwhals._polars.expr import PolarsExpr
 
+        native_dtypes = [
+            narwhals_to_native_dtype(dtype, self._version).__class__
+            if isinstance(dtype, type) and issubclass(dtype, DType)
+            else narwhals_to_native_dtype(dtype, self._version)
+            for dtype in dtypes
+        ]
         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
 
@@ -15,6 +16,7 @@
 from tests.utils import PYARROW_VERSION
 from tests.utils import Constructor
 from tests.utils import assert_equal_data
+from tests.utils import is_windows
 
 data = {
     "a": [1, 1, 2],
@@ -125,3 +127,27 @@ def test_set_ops_invalid(constructor: Constructor) -> None:
         match=re.escape("unsupported operand type(s) for op: ('Selector' + 'Selector')"),
     ):
         df.select(boolean() + numeric())
+
+
+@pytest.mark.skipif(is_windows(), reason="windows is what it is")
+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)
+    if "pyarrow_table" in str(constructor) and PYARROW_VERSION < (12,):
+        # bug in old pyarrow
+        request.applymarker(pytest.mark.xfail)
+    if "duckdb" in str(constructor) or "pyspark" in str(constructor):
+        # replace_time_zone not implemented
+        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
+    result = df.select(nw.selectors.by_dtype(nw.Int64())).collect_schema().names()
+    expected = ["c"]
+    assert result == expected