diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml index bf4f7d39ab..20058a435d 100644 --- a/.github/workflows/pytest.yml +++ b/.github/workflows/pytest.yml @@ -31,6 +31,7 @@ jobs: - name: Run pytest run: pytest tests --cov=narwhals --cov=tests --cov-fail-under=85 - name: Run doctests + if: startsWith(matrix.os, 'windows') != true run: pytest narwhals --doctest-modules pytest-windows: diff --git a/docs/api-reference/expr_dt.md b/docs/api-reference/expr_dt.md index e04e928890..5c9ab41f3c 100644 --- a/docs/api-reference/expr_dt.md +++ b/docs/api-reference/expr_dt.md @@ -4,6 +4,7 @@ handler: python options: members: + - convert_time_zone - date - year - month @@ -15,6 +16,7 @@ - millisecond - microsecond - nanosecond + - replace_time_zone - total_minutes - total_seconds - total_milliseconds diff --git a/docs/api-reference/series_dt.md b/docs/api-reference/series_dt.md index ba342ad302..c925924116 100644 --- a/docs/api-reference/series_dt.md +++ b/docs/api-reference/series_dt.md @@ -4,6 +4,7 @@ handler: python options: members: + - convert_time_zone - date - year - month @@ -15,6 +16,7 @@ - millisecond - microsecond - nanosecond + - replace_time_zone - total_minutes - total_seconds - total_milliseconds diff --git a/narwhals/_arrow/expr.py b/narwhals/_arrow/expr.py index c70425efe3..55c529d304 100644 --- a/narwhals/_arrow/expr.py +++ b/narwhals/_arrow/expr.py @@ -410,6 +410,16 @@ def to_string(self: Self, format: str) -> ArrowExpr: # noqa: A002 self._expr, "dt", "to_string", format ) + def replace_time_zone(self: Self, time_zone: str | None) -> ArrowExpr: + return reuse_series_namespace_implementation( + self._expr, "dt", "replace_time_zone", time_zone + ) + + def convert_time_zone(self: Self, time_zone: str) -> ArrowExpr: + return reuse_series_namespace_implementation( + self._expr, "dt", "convert_time_zone", time_zone + ) + def date(self: Self) -> ArrowExpr: return reuse_series_namespace_implementation(self._expr, "dt", "date") diff --git a/narwhals/_arrow/series.py b/narwhals/_arrow/series.py index 5070234985..65a393ca99 100644 --- a/narwhals/_arrow/series.py +++ b/narwhals/_arrow/series.py @@ -755,6 +755,31 @@ def to_string(self: Self, format: str) -> ArrowSeries: # noqa: A002 pc.strftime(self._arrow_series._native_series, format) ) + def replace_time_zone(self: Self, time_zone: str | None) -> ArrowSeries: + import pyarrow.compute as pc # ignore-banned-import() + + if time_zone is not None: + result = pc.assume_timezone( + pc.local_timestamp(self._arrow_series._native_series), time_zone + ) + else: + result = pc.local_timestamp(self._arrow_series._native_series) + return self._arrow_series._from_native_series(result) + + def convert_time_zone(self: Self, time_zone: str) -> ArrowSeries: + import pyarrow as pa # ignore-banned-import + + if self._arrow_series.dtype.time_zone is None: # type: ignore[attr-defined] + result = self.replace_time_zone("UTC")._native_series.cast( + pa.timestamp(self._arrow_series._native_series.type.unit, time_zone) + ) + else: + result = self._arrow_series._native_series.cast( + pa.timestamp(self._arrow_series._native_series.type.unit, time_zone) + ) + + return self._arrow_series._from_native_series(result) + def date(self: Self) -> ArrowSeries: import pyarrow as pa # ignore-banned-import() diff --git a/narwhals/_dask/expr.py b/narwhals/_dask/expr.py index 10b95bc89e..693fcad5e4 100644 --- a/narwhals/_dask/expr.py +++ b/narwhals/_dask/expr.py @@ -10,6 +10,7 @@ from narwhals._dask.utils import add_row_index from narwhals._dask.utils import maybe_evaluate from narwhals._dask.utils import narwhals_to_native_dtype +from narwhals._pandas_like.utils import native_to_narwhals_dtype from narwhals.utils import generate_unique_token if TYPE_CHECKING: @@ -925,6 +926,33 @@ def to_string(self, format: str) -> DaskExpr: # noqa: A002 returns_scalar=False, ) + def replace_time_zone(self, time_zone: str | None) -> DaskExpr: + return self._expr._from_call( + lambda _input, _time_zone: _input.dt.tz_localize(None).dt.tz_localize( + _time_zone + ) + if _time_zone is not None + else _input.dt.tz_localize(None), + "tz_localize", + time_zone, + returns_scalar=False, + ) + + def convert_time_zone(self, time_zone: str) -> DaskExpr: + def func(s: dask_expr.Series, time_zone: str) -> dask_expr.Series: + dtype = native_to_narwhals_dtype(s, self._expr._dtypes) + if dtype.time_zone is None: # type: ignore[attr-defined] + return s.dt.tz_localize("UTC").dt.tz_convert(time_zone) + else: + return s.dt.tz_convert(time_zone) + + return self._expr._from_call( + func, + "tz_convert", + time_zone, + returns_scalar=False, + ) + def total_minutes(self) -> DaskExpr: return self._expr._from_call( lambda _input: _input.dt.total_seconds() // 60, diff --git a/narwhals/_pandas_like/expr.py b/narwhals/_pandas_like/expr.py index 2ebadbe163..07ba3e56dd 100644 --- a/narwhals/_pandas_like/expr.py +++ b/narwhals/_pandas_like/expr.py @@ -572,6 +572,16 @@ def to_string(self, format: str) -> PandasLikeExpr: # noqa: A002 self._expr, "dt", "to_string", format ) + def replace_time_zone(self, time_zone: str | None) -> PandasLikeExpr: + return reuse_series_namespace_implementation( + self._expr, "dt", "replace_time_zone", time_zone + ) + + def convert_time_zone(self, time_zone: str) -> PandasLikeExpr: + return reuse_series_namespace_implementation( + self._expr, "dt", "convert_time_zone", time_zone + ) + class PandasLikeExprNameNamespace: def __init__(self: Self, expr: PandasLikeExpr) -> None: diff --git a/narwhals/_pandas_like/series.py b/narwhals/_pandas_like/series.py index 9cca664057..2532aea8f4 100644 --- a/narwhals/_pandas_like/series.py +++ b/narwhals/_pandas_like/series.py @@ -928,3 +928,21 @@ def to_string(self, format: str) -> PandasLikeSeries: # noqa: A002 return self._pandas_series._from_native_series( self._pandas_series._native_series.dt.strftime(format) ) + + def replace_time_zone(self, time_zone: str | None) -> PandasLikeSeries: + if time_zone is not None: + result = self._pandas_series._native_series.dt.tz_localize( + None + ).dt.tz_localize(time_zone) + else: + result = self._pandas_series._native_series.dt.tz_localize(None) + return self._pandas_series._from_native_series(result) + + def convert_time_zone(self, time_zone: str) -> PandasLikeSeries: + if self._pandas_series.dtype.time_zone is None: # type: ignore[attr-defined] + result = self._pandas_series._native_series.dt.tz_localize( + "UTC" + ).dt.tz_convert(time_zone) + else: + result = self._pandas_series._native_series.dt.tz_convert(time_zone) + return self._pandas_series._from_native_series(result) diff --git a/narwhals/_pandas_like/utils.py b/narwhals/_pandas_like/utils.py index d7ecc98f20..381a78c8da 100644 --- a/narwhals/_pandas_like/utils.py +++ b/narwhals/_pandas_like/utils.py @@ -218,8 +218,8 @@ def set_axis( return obj.set_axis(index, axis=0, **kwargs) # type: ignore[attr-defined, no-any-return] -def native_to_narwhals_dtype(column: Any, dtypes: DTypes) -> DType: - dtype = str(column.dtype) +def native_to_narwhals_dtype(native_column: Any, dtypes: DTypes) -> DType: + dtype = str(native_column.dtype) pd_datetime_rgx = ( r"^datetime64\[(?Ps|ms|us|ns)(?:, (?P[a-zA-Z\/]+))?\]$" @@ -282,26 +282,30 @@ def native_to_narwhals_dtype(column: Any, dtypes: DTypes) -> DType: return dtypes.Date() if dtype.startswith(("large_list", "list")): return dtypes.List( - arrow_native_to_narwhals_dtype(column.dtype.pyarrow_dtype.value_type, dtypes) + arrow_native_to_narwhals_dtype( + native_column.dtype.pyarrow_dtype.value_type, dtypes + ) ) if dtype.startswith("fixed_size_list"): return dtypes.Array( - arrow_native_to_narwhals_dtype(column.dtype.pyarrow_dtype.value_type, dtypes), - column.dtype.pyarrow_dtype.list_size, + arrow_native_to_narwhals_dtype( + native_column.dtype.pyarrow_dtype.value_type, dtypes + ), + native_column.dtype.pyarrow_dtype.list_size, ) if dtype.startswith("struct"): return dtypes.Struct() if dtype == "object": if ( # pragma: no cover TODO(unassigned): why does this show as uncovered? - idx := getattr(column, "first_valid_index", lambda: None)() - ) is not None and isinstance(column.loc[idx], str): + idx := getattr(native_column, "first_valid_index", lambda: None)() + ) is not None and isinstance(native_column.loc[idx], str): # Infer based on first non-missing value. # For pandas pre 3.0, this isn't perfect. # After pandas 3.0, pandas has a dedicated string dtype # which is inferred by default. return dtypes.String() else: - df = column.to_frame() + df = native_column.to_frame() if hasattr(df, "__dataframe__"): from narwhals._interchange.dataframe import ( map_interchange_dtype_to_narwhals_dtype, diff --git a/narwhals/dataframe.py b/narwhals/dataframe.py index e4ad31b388..b0ac1c329f 100644 --- a/narwhals/dataframe.py +++ b/narwhals/dataframe.py @@ -7,6 +7,7 @@ from typing import Iterable from typing import Iterator from typing import Literal +from typing import NoReturn from typing import Sequence from typing import TypeVar from typing import overload @@ -2787,7 +2788,7 @@ def __repr__(self) -> str: # pragma: no cover + "┘" ) - def __getitem__(self, item: str | slice) -> Series | Self: + def __getitem__(self, item: str | slice) -> NoReturn: msg = "Slicing is not supported on LazyFrame" raise TypeError(msg) diff --git a/narwhals/expr.py b/narwhals/expr.py index 8446d81c38..abf9c10430 100644 --- a/narwhals/expr.py +++ b/narwhals/expr.py @@ -3501,6 +3501,119 @@ def to_string(self, format: str) -> Expr: # noqa: A002 lambda plx: self._expr._call(plx).dt.to_string(format) ) + def replace_time_zone(self, time_zone: str | None) -> Expr: + """ + Replace time zone. + + Arguments: + time_zone: Target time zone. + + Examples: + >>> from datetime import datetime, timezone + >>> import narwhals as nw + >>> import pandas as pd + >>> import polars as pl + >>> import pyarrow as pa + >>> data = { + ... "a": [ + ... datetime(2024, 1, 1, tzinfo=timezone.utc), + ... datetime(2024, 1, 2, tzinfo=timezone.utc), + ... ] + ... } + >>> df_pd = pd.DataFrame(data) + >>> df_pl = pl.DataFrame(data) + >>> df_pa = pa.table(data) + + Let's define a dataframe-agnostic function: + + >>> @nw.narwhalify + ... def func(df): + ... return df.select(nw.col("a").dt.replace_time_zone("Asia/Kathmandu")) + + We can then pass pandas / PyArrow / Polars / any other supported library: + + >>> func(df_pd) + a + 0 2024-01-01 00:00:00+05:45 + 1 2024-01-02 00:00:00+05:45 + >>> func(df_pl) + shape: (2, 1) + ┌──────────────────────────────┐ + │ a │ + │ --- │ + │ datetime[μs, Asia/Kathmandu] │ + ╞══════════════════════════════╡ + │ 2024-01-01 00:00:00 +0545 │ + │ 2024-01-02 00:00:00 +0545 │ + └──────────────────────────────┘ + >>> func(df_pa) # doctest:+SKIP + pyarrow.Table + a: timestamp[us, tz=Asia/Kathmandu] + ---- + a: [[2023-12-31 18:15:00.000000Z,2024-01-01 18:15:00.000000Z]] + """ + return self._expr.__class__( + lambda plx: self._expr._call(plx).dt.replace_time_zone(time_zone) + ) + + def convert_time_zone(self, time_zone: str) -> Expr: + """ + Convert to a new time zone. + + Arguments: + time_zone: Target time zone. + + Examples: + >>> from datetime import datetime, timezone + >>> import narwhals as nw + >>> import pandas as pd + >>> import polars as pl + >>> import pyarrow as pa + >>> data = { + ... "a": [ + ... datetime(2024, 1, 1, tzinfo=timezone.utc), + ... datetime(2024, 1, 2, tzinfo=timezone.utc), + ... ] + ... } + >>> df_pd = pd.DataFrame(data) + >>> df_pl = pl.DataFrame(data) + >>> df_pa = pa.table(data) + + Let's define a dataframe-agnostic function: + + >>> @nw.narwhalify + ... def func(df): + ... return df.select(nw.col("a").dt.convert_time_zone("Asia/Kathmandu")) + + We can then pass pandas / PyArrow / Polars / any other supported library: + + >>> func(df_pd) + a + 0 2024-01-01 05:45:00+05:45 + 1 2024-01-02 05:45:00+05:45 + >>> func(df_pl) + shape: (2, 1) + ┌──────────────────────────────┐ + │ a │ + │ --- │ + │ datetime[μs, Asia/Kathmandu] │ + ╞══════════════════════════════╡ + │ 2024-01-01 05:45:00 +0545 │ + │ 2024-01-02 05:45:00 +0545 │ + └──────────────────────────────┘ + >>> func(df_pa) # doctest:+SKIP + pyarrow.Table + a: timestamp[us, tz=Asia/Kathmandu] + ---- + a: [[2024-01-01 00:00:00.000000Z,2024-01-02 00:00:00.000000Z]] + """ + if time_zone is None: + msg = "Target `time_zone` cannot be `None` in `convert_time_zone`. Please use `replace_time_zone(None)` if you want to remove the time zone." + raise TypeError(msg) + return self._expr.__class__( + lambda plx: self._expr._call(plx).dt.convert_time_zone(time_zone) + ) + class ExprNameNamespace: def __init__(self: Self, expr: Expr) -> None: diff --git a/narwhals/series.py b/narwhals/series.py index 1753598c17..a9eb51a421 100644 --- a/narwhals/series.py +++ b/narwhals/series.py @@ -3880,3 +3880,112 @@ def to_string(self, format: str) -> Series: # noqa: A002 return self._narwhals_series._from_compliant_series( self._narwhals_series._compliant_series.dt.to_string(format) ) + + def replace_time_zone(self, time_zone: str | None) -> Series: + """ + Replace time zone. + + Arguments: + time_zone: Target time zone. + + Examples: + >>> from datetime import datetime, timezone + >>> import narwhals as nw + >>> import pandas as pd + >>> import polars as pl + >>> import pyarrow as pa + >>> data = [ + ... datetime(2024, 1, 1, tzinfo=timezone.utc), + ... datetime(2024, 1, 2, tzinfo=timezone.utc), + ... ] + >>> s_pd = pd.Series(data) + >>> s_pl = pl.Series(data) + >>> s_pa = pa.chunked_array([data]) + + Let's define a dataframe-agnostic function: + + >>> @nw.narwhalify + ... def func(s): + ... return s.dt.replace_time_zone("Asia/Kathmandu") + + We can then pass pandas / PyArrow / Polars / any other supported library: + + >>> func(s_pd) + 0 2024-01-01 00:00:00+05:45 + 1 2024-01-02 00:00:00+05:45 + dtype: datetime64[ns, Asia/Kathmandu] + >>> func(s_pl) # doctest: +NORMALIZE_WHITESPACE + shape: (2,) + Series: '' [datetime[μs, Asia/Kathmandu]] + [ + 2024-01-01 00:00:00 +0545 + 2024-01-02 00:00:00 +0545 + ] + >>> func(s_pa) # doctest: +SKIP + + [ + [ + 2023-12-31 18:15:00.000000Z, + 2024-01-01 18:15:00.000000Z + ] + ] + """ + return self._narwhals_series._from_compliant_series( + self._narwhals_series._compliant_series.dt.replace_time_zone(time_zone) + ) + + def convert_time_zone(self, time_zone: str) -> Series: + """ + Convert time zone. + + Arguments: + time_zone: Target time zone. + + Examples: + >>> from datetime import datetime, timezone + >>> import narwhals as nw + >>> import pandas as pd + >>> import polars as pl + >>> import pyarrow as pa + >>> data = [ + ... datetime(2024, 1, 1, tzinfo=timezone.utc), + ... datetime(2024, 1, 2, tzinfo=timezone.utc), + ... ] + >>> s_pd = pd.Series(data) + >>> s_pl = pl.Series(data) + >>> s_pa = pa.chunked_array([data]) + + Let's define a dataframe-agnostic function: + + >>> @nw.narwhalify + ... def func(s): + ... return s.dt.convert_time_zone("Asia/Kathmandu") + + We can then pass pandas / PyArrow / Polars / any other supported library: + + >>> func(s_pd) + 0 2024-01-01 05:45:00+05:45 + 1 2024-01-02 05:45:00+05:45 + dtype: datetime64[ns, Asia/Kathmandu] + >>> func(s_pl) # doctest: +NORMALIZE_WHITESPACE + shape: (2,) + Series: '' [datetime[μs, Asia/Kathmandu]] + [ + 2024-01-01 05:45:00 +0545 + 2024-01-02 05:45:00 +0545 + ] + >>> func(s_pa) # doctest: +SKIP + + [ + [ + 2024-01-01 00:00:00.000000Z, + 2024-01-02 00:00:00.000000Z + ] + ] + """ + if time_zone is None: + msg = "Target `time_zone` cannot be `None` in `convert_time_zone`. Please use `replace_time_zone(None)` if you want to remove the time zone." + raise TypeError(msg) + return self._narwhals_series._from_compliant_series( + self._narwhals_series._compliant_series.dt.convert_time_zone(time_zone) + ) diff --git a/tests/expr_and_series/convert_time_zone_test.py b/tests/expr_and_series/convert_time_zone_test.py new file mode 100644 index 0000000000..bc5d176bad --- /dev/null +++ b/tests/expr_and_series/convert_time_zone_test.py @@ -0,0 +1,120 @@ +from datetime import datetime +from datetime import timezone +from typing import Any + +import pandas as pd +import polars as pl +import pyarrow as pa +import pytest + +import narwhals.stable.v1 as nw +from narwhals.utils import parse_version +from tests.utils import Constructor +from tests.utils import compare_dicts +from tests.utils import is_windows + + +def test_convert_time_zone( + constructor: Constructor, request: pytest.FixtureRequest +) -> None: + if (any(x in str(constructor) for x in ("pyarrow", "modin")) and is_windows()) or ( + "pandas_pyarrow" in str(constructor) and parse_version(pd.__version__) < (2, 2) + ): + request.applymarker(pytest.mark.xfail) + data = { + "a": [ + datetime(2020, 1, 1, tzinfo=timezone.utc), + datetime(2020, 1, 2, tzinfo=timezone.utc), + ] + } + df = nw.from_native(constructor(data)) + result = df.select(nw.col("a").dt.convert_time_zone("Asia/Kathmandu")) + result_dtype = result.collect_schema()["a"] + assert result_dtype == nw.Datetime + assert result_dtype.time_zone == "Asia/Kathmandu" # type: ignore[attr-defined] + result_str = result.select(nw.col("a").dt.to_string("%Y-%m-%dT%H:%M%z")) + expected = {"a": ["2020-01-01T05:45+0545", "2020-01-02T05:45+0545"]} + compare_dicts(result_str, expected) + + +def test_convert_time_zone_series( + constructor_eager: Any, request: pytest.FixtureRequest +) -> None: + if ( + any(x in str(constructor_eager) for x in ("pyarrow", "modin")) and is_windows() + ) or ( + "pandas_pyarrow" in str(constructor_eager) + and parse_version(pd.__version__) < (2, 2) + ): + request.applymarker(pytest.mark.xfail) + data = { + "a": [ + datetime(2020, 1, 1, tzinfo=timezone.utc), + datetime(2020, 1, 2, tzinfo=timezone.utc), + ] + } + df = nw.from_native(constructor_eager(data), eager_only=True) + result = df.select(df["a"].dt.convert_time_zone("Asia/Kathmandu")) + result_dtype = result.collect_schema()["a"] + assert result_dtype == nw.Datetime + assert result_dtype.time_zone == "Asia/Kathmandu" # type: ignore[attr-defined] + result_str = result.select(nw.col("a").dt.to_string("%Y-%m-%dT%H:%M%z")) + expected = {"a": ["2020-01-01T05:45+0545", "2020-01-02T05:45+0545"]} + compare_dicts(result_str, expected) + + +def test_convert_time_zone_from_none( + constructor: Constructor, request: pytest.FixtureRequest +) -> None: + if ( + (any(x in str(constructor) for x in ("pyarrow", "modin")) and is_windows()) + or ( + "pandas_pyarrow" in str(constructor) + and parse_version(pd.__version__) < (2, 2) + ) + or ("pyarrow_table" in str(constructor) and parse_version(pa.__version__) < (12,)) + ): + request.applymarker(pytest.mark.xfail) + if "polars" in str(constructor) and parse_version(pl.__version__) < (0, 20, 7): + # polars used to disallow this + request.applymarker(pytest.mark.xfail) + data = { + "a": [ + datetime(2020, 1, 1, tzinfo=timezone.utc), + datetime(2020, 1, 2, tzinfo=timezone.utc), + ] + } + df = nw.from_native(constructor(data)) + result = df.select( + nw.col("a").dt.replace_time_zone(None).dt.convert_time_zone("Asia/Kathmandu") + ) + result_dtype = result.collect_schema()["a"] + assert result_dtype == nw.Datetime + assert result_dtype.time_zone == "Asia/Kathmandu" # type: ignore[attr-defined] + result_str = result.select(nw.col("a").dt.to_string("%Y-%m-%dT%H:%M%z")) + expected = {"a": ["2020-01-01T05:45+0545", "2020-01-02T05:45+0545"]} + compare_dicts(result_str, expected) + + +def test_convert_time_zone_to_none(constructor: Constructor) -> None: + data = { + "a": [ + datetime(2020, 1, 1, tzinfo=timezone.utc), + datetime(2020, 1, 2, tzinfo=timezone.utc), + ] + } + df = nw.from_native(constructor(data)) + with pytest.raises(TypeError, match="Target `time_zone` cannot be `None`"): + df.select(nw.col("a").dt.convert_time_zone(None)) # type: ignore[arg-type] + + +def test_convert_time_zone_to_none_series(constructor_eager: Any) -> None: + data = { + "a": [ + datetime(2020, 1, 1, tzinfo=timezone.utc), + datetime(2020, 1, 2, tzinfo=timezone.utc), + ] + } + df = nw.from_native(constructor_eager(data)) + with pytest.raises(TypeError, match="Target `time_zone` cannot be `None`"): + df["a"].dt.convert_time_zone(None) # type: ignore[arg-type] diff --git a/tests/expr_and_series/replace_time_zone_test.py b/tests/expr_and_series/replace_time_zone_test.py new file mode 100644 index 0000000000..560fcfe84d --- /dev/null +++ b/tests/expr_and_series/replace_time_zone_test.py @@ -0,0 +1,125 @@ +from datetime import datetime +from datetime import timezone +from typing import Any + +import pandas as pd +import pyarrow as pa +import pytest + +import narwhals.stable.v1 as nw +from narwhals.utils import parse_version +from tests.utils import Constructor +from tests.utils import compare_dicts +from tests.utils import is_windows + + +def test_replace_time_zone( + constructor: Constructor, request: pytest.FixtureRequest +) -> None: + if ( + (any(x in str(constructor) for x in ("pyarrow", "modin")) and is_windows()) + or ("pandas_pyarrow" in str(constructor) and parse_version(pd.__version__) < (2,)) + or ("pyarrow_table" in str(constructor) and parse_version(pa.__version__) < (12,)) + ): + request.applymarker(pytest.mark.xfail) + data = { + "a": [ + datetime(2020, 1, 1, tzinfo=timezone.utc), + datetime(2020, 1, 2, tzinfo=timezone.utc), + ] + } + df = nw.from_native(constructor(data)) + result = df.select(nw.col("a").dt.replace_time_zone("Asia/Kathmandu")) + result_dtype = result.collect_schema()["a"] + assert result_dtype == nw.Datetime + assert result_dtype.time_zone == "Asia/Kathmandu" # type: ignore[attr-defined] + result_str = result.select(nw.col("a").dt.to_string("%Y-%m-%dT%H:%M%z")) + expected = {"a": ["2020-01-01T00:00+0545", "2020-01-02T00:00+0545"]} + compare_dicts(result_str, expected) + + +def test_replace_time_zone_none( + constructor: Constructor, request: pytest.FixtureRequest +) -> None: + if ( + (any(x in str(constructor) for x in ("pyarrow", "modin")) and is_windows()) + or ("pandas_pyarrow" in str(constructor) and parse_version(pd.__version__) < (2,)) + or ("pyarrow_table" in str(constructor) and parse_version(pa.__version__) < (12,)) + ): + request.applymarker(pytest.mark.xfail) + data = { + "a": [ + datetime(2020, 1, 1, tzinfo=timezone.utc), + datetime(2020, 1, 2, tzinfo=timezone.utc), + ] + } + df = nw.from_native(constructor(data)) + result = df.select(nw.col("a").dt.replace_time_zone(None)) + result_dtype = result.collect_schema()["a"] + assert result_dtype == nw.Datetime + assert result_dtype.time_zone is None # type: ignore[attr-defined] + result_str = result.select(nw.col("a").dt.to_string("%Y-%m-%dT%H:%M")) + expected = {"a": ["2020-01-01T00:00", "2020-01-02T00:00"]} + compare_dicts(result_str, expected) + + +def test_replace_time_zone_series( + constructor_eager: Any, request: pytest.FixtureRequest +) -> None: + if ( + (any(x in str(constructor_eager) for x in ("pyarrow", "modin")) and is_windows()) + or ( + "pandas_pyarrow" in str(constructor_eager) + and parse_version(pd.__version__) < (2,) + ) + or ( + "pyarrow_table" in str(constructor_eager) + and parse_version(pa.__version__) < (12,) + ) + ): + request.applymarker(pytest.mark.xfail) + data = { + "a": [ + datetime(2020, 1, 1, tzinfo=timezone.utc), + datetime(2020, 1, 2, tzinfo=timezone.utc), + ] + } + df = nw.from_native(constructor_eager(data), eager_only=True) + result = df.select(df["a"].dt.replace_time_zone("Asia/Kathmandu")) + result_dtype = result.collect_schema()["a"] + assert result_dtype == nw.Datetime + assert result_dtype.time_zone == "Asia/Kathmandu" # type: ignore[attr-defined] + result_str = result.select(nw.col("a").dt.to_string("%Y-%m-%dT%H:%M%z")) + expected = {"a": ["2020-01-01T00:00+0545", "2020-01-02T00:00+0545"]} + compare_dicts(result_str, expected) + + +def test_replace_time_zone_none_series( + constructor_eager: Any, request: pytest.FixtureRequest +) -> None: + if ( + (any(x in str(constructor_eager) for x in ("pyarrow", "modin")) and is_windows()) + or ( + "pandas_pyarrow" in str(constructor_eager) + and parse_version(pd.__version__) < (2,) + ) + or ( + "pyarrow_table" in str(constructor_eager) + and parse_version(pa.__version__) < (12,) + ) + ): + request.applymarker(pytest.mark.xfail) + data = { + "a": [ + datetime(2020, 1, 1, tzinfo=timezone.utc), + datetime(2020, 1, 2, tzinfo=timezone.utc), + ] + } + df = nw.from_native(constructor_eager(data), eager_only=True) + result = df.select(df["a"].dt.replace_time_zone(None)) + result_dtype = result.collect_schema()["a"] + assert result_dtype == nw.Datetime + assert result_dtype.time_zone is None # type: ignore[attr-defined] + result_str = result.select(df["a"].dt.to_string("%Y-%m-%dT%H:%M")) + expected = {"a": ["2020-01-01T00:00", "2020-01-02T00:00"]} + compare_dicts(result_str, expected)