Repository: narwhals-dev/narwhals (python)

Issue:
selectors.by_dtype doesn't select tz-aware columns
```python
import polars as pl
import pandas as pd
import narwhals as nw
import pyarrow as pa
import numpy as np
from datetime import datetime

data = {'a': [datetime(2020, 1, 1), datetime(2020, 1, 2)], 'c': [4,5]}

df = nw.from_native(pd.DataFrame(data)).with_columns(b=nw.col('a').dt.replace_time_zone('Asia/Katmandu'))
print(df.select(nw.selectors.by_dtype(nw.Datetime)).to_native())

df = nw.from_native(pl.DataFrame(data)).with_columns(b=nw.col('a').dt.replace_time_zone('Asia/Katmandu'))
print(df.select(nw.selectors.by_dtype(nw.Datetime)).to_native())
```
this outputs
```
           a                         b
0 2020-01-01 2020-01-01 00:00:00+05:45
1 2020-01-02 2020-01-02 00:00:00+05:45
shape: (2, 1)
┌─────────────────────┐
│ a                   │
│ ---                 │
│ datetime[μs]        │
╞═════════════════════╡
│ 2020-01-01 00:00:00 │
│ 2020-01-02 00:00:00 │
└─────────────────────┘
```

For Polars, both columns 'a' and 'b' should be selected

Pull request:
fix: `nw.selectors.by_dtype(nw.Datetime)` was excluding tz-aware for Polars
<!--
# Thanks for contributing a pull request! 
## Please make sure you see our contribution guidelines: https://github.com/narwhals-dev/narwhals/blob/main/CONTRIBUTING.md
-->

closes #1313 

## What type of PR is this? (check all applicable)

- [ ] 💾 Refactor
- [ ] ✨ Feature
- [ ] 🐛 Bug Fix
- [ ] 🔧 Optimization
- [ ] 📝 Documentation
- [ ] ✅ Test
- [ ] 🐳 Other

## Related issues

- Related issue #\<issue number\>
- Closes #\<issue number\>

## Checklist

- [ ] Code follows style guide (ruff)
- [ ] Tests added
- [ ] Documented the changes

## If you have comments or can explain your changes, please do so below


Patch:
diff --git a/narwhals/_polars/namespace.py b/narwhals/_polars/namespace.py
index 32e53b37..005af360 100644
--- a/narwhals/_polars/namespace.py
+++ b/narwhals/_polars/namespace.py
@@ -13,6 +13,7 @@ import polars as pl
 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 @@ if TYPE_CHECKING:
     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 @@ class PolarsNamespace:
         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 @@ class PolarsNamespace:
         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 @@ class PolarsNamespace:
         )
 
     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 @@ class PolarsNamespace:
         )
 
     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 @@ class PolarsNamespace:
         )
 
     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 @@ class PolarsNamespace:
         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 @@ class PolarsSelectors:
         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,
         )


Test patch:
diff --git a/tests/selectors_test.py b/tests/selectors_test.py
index 99740d22..9e16da5a 100644
--- a/tests/selectors_test.py
+++ b/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 POLARS_VERSION
 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