sort out dataframe.to_numpy too

narwhals-dev · Nov 2, 2024 · 461244b · 461244b
1 parent 75badf8
commit 461244b
Showing 1 changed file with 3 additions and 6 deletions.
diff --git a/narwhals/_pandas_like/dataframe.py b/narwhals/_pandas_like/dataframe.py
@@ -701,17 +701,14 @@ def to_numpy(self, dtype: Any = None, copy: bool | None = None) -> Any:
         if dtype is not None:
             return self._native_frame.to_numpy(dtype=dtype, copy=copy)
 
-        convert = set()
-        for key, val in self.schema.items():
-            if val == self._dtypes.Datetime and val.time_zone is not None:  # type: ignore[attr-defined]
-                convert.add(key)
         df = self.with_columns(
             *[
                 self.__narwhals_namespace__()
-                .col(x)
+                .col(key)
                 .dt.convert_time_zone("UTC")
                 .dt.replace_time_zone(None)
-                for x in convert
+                for key, val in self.schema.items()
+                if val == self._dtypes.Datetime and val.time_zone is not None  # type: ignore[attr-defined]
             ]
         )._native_frame