feat: Datetime(time_unit, time_zone) and Duration(time_unit) types

narwhals/_arrow/utils.py

@@ -46,7 +46,7 @@ def translate_dtype(dtype: Any) -> dtypes.DType:
     if pa.types.is_date32(dtype):
         return dtypes.Date()
     if pa.types.is_timestamp(dtype):
-        return dtypes.Datetime()
+        return dtypes.Datetime(time_unit=dtype.unit, time_zone=dtype.tz)
     if pa.types.is_duration(dtype):
         return dtypes.Duration()
     if pa.types.is_dictionary(dtype):
@@ -88,8 +88,10 @@ def narwhals_to_native_dtype(dtype: dtypes.DType | type[dtypes.DType]) -> Any:
         # with Polars for now
         return pa.dictionary(pa.uint32(), pa.string())
     if isinstance_or_issubclass(dtype, dtypes.Datetime):
-        # Use Polars' default
-        return pa.timestamp("us")
+        time_unit = getattr(dtype, "time_unit", "us")
+        time_zone = getattr(dtype, "time_zone", None)
+        return pa.timestamp(time_unit, tz=time_zone)
+
     if isinstance_or_issubclass(dtype, dtypes.Duration):
         # Use Polars' default
         return pa.duration("us")

narwhals/_pandas_like/utils.py

@@ -1,8 +1,10 @@
 from __future__ import annotations
 
+import re
 from typing import TYPE_CHECKING
 from typing import Any
 from typing import Iterable
+from typing import Literal
 from typing import TypeVar
 
 from narwhals.dependencies import get_cudf
@@ -221,6 +223,12 @@ def translate_dtype(column: Any) -> DType:
     from narwhals import dtypes
 
     dtype = column.dtype
+
+    pd_datetime_rgx = (
+        r"^datetime64\[(?P<time_unit>ms|us|ns)(?:, (?P<time_zone>[a-zA-Z\/]+))?\]$"
+    )
+    pa_datetime_rgx = r"^timestamp\[(?P<time_unit>ms|us|ns)(?:, tz=(?P<time_zone>[a-zA-Z\/]+))?\]\[pyarrow\]$"
+
     if str(dtype) in ("int64", "Int64", "Int64[pyarrow]", "int64[pyarrow]"):
         return dtypes.Int64()
     if str(dtype) in ("int32", "Int32", "Int32[pyarrow]", "int32[pyarrow]"):
@@ -264,16 +272,15 @@ def translate_dtype(column: Any) -> DType:
         return dtypes.Boolean()
     if str(dtype) in ("category",) or str(dtype).startswith("dictionary<"):
         return dtypes.Categorical()
-    if str(dtype).startswith("datetime64"):
-        # TODO(Unassigned): different time units and time zones
-        return dtypes.Datetime()
+    if (match_ := re.match(pd_datetime_rgx, str(dtype))) or (
+        match_ := re.match(pa_datetime_rgx, str(dtype))
+    ):
+        time_unit: Literal["us", "ns", "ms"] = match_.group("time_unit")  # type: ignore[assignment]
+        time_zone: str | None = match_.group("time_zone")
+        return dtypes.Datetime(time_unit, time_zone)
     if str(dtype).startswith("timedelta64") or str(dtype).startswith("duration"):
         # TODO(Unassigned): different time units
         return dtypes.Duration()
-    if str(dtype).startswith("timestamp["):
-        # pyarrow-backed datetime
-        # TODO(Unassigned): different time units and time zones
-        return dtypes.Datetime()
     if str(dtype) == "date32[day][pyarrow]":
         return dtypes.Date()
     if str(dtype) == "object":
@@ -425,10 +432,16 @@ def narwhals_to_native_dtype(  # noqa: PLR0915
         # convert to it?
         return "category"
     if isinstance_or_issubclass(dtype, dtypes.Datetime):
-        # TODO(Unassigned): different time units and time zones
+        time_unit = getattr(dtype, "time_unit", "us")
+        time_zone = getattr(dtype, "time_zone", None)
+
         if dtype_backend == "pyarrow-nullable":
-            return "timestamp[ns][pyarrow]"
-        return "datetime64[ns]"
+            tz_part = f", tz={time_zone}" if time_zone else ""
+            return f"timestamp[{time_unit}{tz_part}][pyarrow]"
+        else:
+            tz_part = f", {time_zone}" if time_zone else ""
+            return f"datetime64[{time_unit}{tz_part}]"
+
     if isinstance_or_issubclass(dtype, dtypes.Duration):
         # TODO(Unassigned): different time units and time zones
         if dtype_backend == "pyarrow-nullable":

narwhals/_polars/utils.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from typing import Any
+from typing import Literal
 
 from narwhals import dtypes
 from narwhals.dependencies import get_polars
@@ -59,12 +60,14 @@ def translate_dtype(dtype: Any) -> dtypes.DType:
         return dtypes.Categorical()
     if dtype == pl.Enum:
         return dtypes.Enum()
-    if dtype == pl.Datetime:
-        return dtypes.Datetime()
     if dtype == pl.Duration:
         return dtypes.Duration()
     if dtype == pl.Date:
         return dtypes.Date()
+    if dtype == pl.Datetime or isinstance(dtype, pl.Datetime):
+        time_unit: Literal["us", "ns", "ms"] = getattr(dtype, "time_unit", "us")
+        time_zone = getattr(dtype, "time_zone", None)
+        return dtypes.Datetime(time_unit=time_unit, time_zone=time_zone)
     return dtypes.Unknown()
 
 
@@ -103,10 +106,13 @@ def narwhals_to_native_dtype(dtype: dtypes.DType | type[dtypes.DType]) -> Any:
     if dtype == dtypes.Enum:
         msg = "Converting to Enum is not (yet) supported"
         raise NotImplementedError(msg)
-    if dtype == dtypes.Datetime:
-        return pl.Datetime()
     if dtype == dtypes.Duration:
         return pl.Duration()
     if dtype == dtypes.Date:
         return pl.Date()
+    if dtype == dtypes.Datetime or isinstance(dtype, dtypes.Datetime):
+        time_unit = getattr(dtype, "time_unit", "us")
+        time_zone = getattr(dtype, "time_zone", None)
+        return pl.Datetime(time_unit, time_zone)
+
     return pl.Unknown()  # pragma: no cover

narwhals/dtypes.py

@@ -1,6 +1,8 @@
 from __future__ import annotations
 
+from datetime import timezone
 from typing import TYPE_CHECKING
+from typing import Literal
 
 if TYPE_CHECKING:
     from typing_extensions import Self
@@ -71,7 +73,39 @@ class Object(DType): ...
 class Unknown(DType): ...
 
 
-class Datetime(TemporalType): ...
+class Datetime(TemporalType):
+    """
+    Data type representing a calendar date and time of day.
+
+    Arguments:
+        time_unit: Unit of time. Defaults to `'us'` (microseconds).
+        time_zone: Time zone string, as defined in zoneinfo (to see valid strings run
+            `import zoneinfo; zoneinfo.available_timezones()` for a full list).
+            When used to match dtypes, can set this to "*" to check for Datetime
+            columns that have any (non-null) timezone.
+
+    Notes:
+        Adapted from Polars implementation at:
+        https://github.com/pola-rs/polars/blob/py-1.7.1/py-polars/polars/datatypes/classes.py#L398-L457
+    """
+
+    def __init__(
+        self: Self,
+        time_unit: Literal["us", "ns", "ms"] = "us",
+        time_zone: str | timezone | None = None,
+    ) -> None:
+        if time_unit not in {"ms", "us", "ns"}:
+            msg = (
+                "invalid `time_unit`"
+                f"\n\nExpected one of {{'ns','us','ms'}}, got {time_unit!r}."
+            )
+            raise ValueError(msg)
+
+        if isinstance(time_zone, timezone):
+            time_zone = str(time_zone)
+
+        self.time_unit = time_unit
+        self.time_zone = time_zone
 
 
 class Duration(TemporalType): ...

tests/dtypes_test.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+from datetime import timezone
+from types import NoneType
+from typing import Literal
+
+import pytest
+
+import narwhals.stable.v1 as nw
+
+
+@pytest.mark.parametrize("time_unit", ["us", "ns", "ms"])
+@pytest.mark.parametrize("time_zone", ["Europe/Rome", timezone.utc, None])
+def test_datetime_valid(
+    time_unit: Literal["us", "ns", "ms"], time_zone: str | timezone | None
+) -> None:
+    dtype = nw.Datetime(time_unit=time_unit, time_zone=time_zone)
+
+    assert dtype.time_unit == time_unit
+    assert isinstance(dtype.time_zone, (str, NoneType))
+
+
+@pytest.mark.parametrize("time_unit", ["abc", "s"])
+def test_datetime_invalid(time_unit: str) -> None:
+    with pytest.raises(ValueError, match="invalid `time_unit`"):
+        nw.Datetime(time_unit=time_unit)  # type: ignore[arg-type]

tests/expr_and_series/cast_test.py

@@ -1,3 +1,8 @@
+from __future__ import annotations
+
+from datetime import datetime
+from datetime import timedelta
+from datetime import timezone
 from typing import Any
 
 import pandas as pd
@@ -6,6 +11,7 @@
 
 import narwhals.stable.v1 as nw
 from narwhals.utils import parse_version
+from tests.utils import compare_dicts
 
 data = {
     "a": [1],
@@ -179,3 +185,27 @@ class Banana:
 
     with pytest.raises(AssertionError, match=r"Unknown dtype"):
         df.select(nw.col("a").cast(Banana))
+
+
+def test_cast_datetime_tz_aware(constructor: Any, request: Any) -> None:
+    if "dask" in str(constructor):
+        request.applymarker(pytest.mark.xfail)
+
+    data = {
+        "date": [
+            datetime(2024, 1, 1, tzinfo=timezone.utc) + timedelta(days=i)
+            for i in range(3)
+        ]
+    }
+    expected = {
+        "date": ["2024-01-01 01:00:00", "2024-01-02 01:00:00", "2024-01-03 01:00:00"]
+    }
+
+    df = nw.from_native(constructor(data))
+    result = df.select(
+        nw.col("date")
+        .cast(nw.Datetime("ms", time_zone="Europe/Rome"))
+        .cast(nw.String())
+        .str.slice(offset=0, length=19)
+    )
+    compare_dicts(result, expected)

tests/series_only/cast_test.py

@@ -75,13 +75,13 @@ def test_cast_date_datetime_pandas() -> None:
     df = df.select(nw.col("a").cast(nw.Datetime))
     result = nw.to_native(df)
     expected = pd.DataFrame({"a": [datetime(2020, 1, 1), datetime(2020, 1, 2)]}).astype(
-        {"a": "timestamp[ns][pyarrow]"}
+        {"a": "timestamp[us][pyarrow]"}
     )
     pd.testing.assert_frame_equal(result, expected)
 
-    # pandas: pyarrow datetime to date
+    # # pandas: pyarrow datetime to date
     dfpd = pd.DataFrame({"a": [datetime(2020, 1, 1), datetime(2020, 1, 2)]}).astype(
-        {"a": "timestamp[ns][pyarrow]"}
+        {"a": "timestamp[us][pyarrow]"}
     )
     df = nw.from_native(dfpd)
     df = df.select(nw.col("a").cast(nw.Date))