From b393e4282a106ab57099197e5df777e3403323ff Mon Sep 17 00:00:00 2001 From: Marco Edward Gorelli Date: Sat, 2 Nov 2024 13:13:05 +0000 Subject: [PATCH] Apply suggestions from code review Co-authored-by: Francesco Bruzzesi <42817048+FBruzzesi@users.noreply.github.com> --- narwhals/_pandas_like/dataframe.py | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/narwhals/_pandas_like/dataframe.py b/narwhals/_pandas_like/dataframe.py index f5497f728..1c90eca79 100644 --- a/narwhals/_pandas_like/dataframe.py +++ b/narwhals/_pandas_like/dataframe.py @@ -698,19 +698,16 @@ def to_numpy(self, dtype: Any = None, copy: bool | None = None) -> Any: # pandas default differs from Polars, but cuDF default is True copy = self._implementation is Implementation.CUDF - to_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] - to_convert.add(key) + to_convert = { + key for key, val in self.schema.items() + if val == self._dtypes.Datetime and val.time_zone is not None # type: ignore[attr-defined] + } if to_convert: df = self.with_columns( - *[ self.__narwhals_namespace__() - .col(key) + .col(list(to_convert)) .dt.convert_time_zone("UTC") .dt.replace_time_zone(None) - for key in to_convert - ] )._native_frame else: df = self._native_frame