python: improve date/datetime get ergonomics

pola-rs · Jan 7, 2022 · 941c49a · 941c49a
1 parent b02bf40
commit 941c49a
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/py-polars/polars/internals/series.py b/py-polars/polars/internals/series.py
@@ -413,7 +413,7 @@ def __getitem__(self, item: Union[int, "Series", range, slice]) -> Any:
         if isinstance(item, int):
             if item < 0:
                 item = self.len() + item
-            if self.dtype in (PlList, Date, Datetime, Duration, Object):
+            if self.dtype in (PlList, Object):
                 f = get_ffi_func("get_<>", self.dtype, self._s)
                 if f is None:
                     return NotImplemented
@@ -3872,8 +3872,7 @@ def truncate(
 
     def __getitem__(self, item: int) -> Union[date, datetime]:
         s = wrap_s(self._s)
-        out = s[item]
-        return _to_python_datetime(out, s.dtype, s.time_unit)
+        return s[item]
 
     def strftime(self, fmt: str) -> Series:
         """

diff --git a/py-polars/tests/test_interop.py b/py-polars/tests/test_interop.py
@@ -21,8 +21,8 @@ def test_from_pandas_datetime() -> None:
         "2021-06-24 00:00:00", "2021-06-24 10:00:00", freq="1H", closed="left"
     )
     s = pl.from_pandas(date_times)
-    assert s[0] == 1624492800000
-    assert s[-1] == 1624525200000
+    assert s[0] == datetime.datetime(2021, 6, 24, 0, 0)
+    assert s[-1] == datetime.datetime(2021, 6, 24, 9, 0)
 
     df = pd.DataFrame({"datetime": ["2021-01-01", "2021-01-02"], "foo": [1, 2]})
     df["datetime"] = pd.to_datetime(df["datetime"])