test_timedeltas.py

from datetime import timedelta

import numpy as np
import pytest

from pandas._libs.tslibs.dtypes import NpyDatetimeUnit

import pandas as pd
from pandas import Timedelta
import pandas._testing as tm
from pandas.core.arrays import (
    DatetimeArray,
    TimedeltaArray,
)


class TestNonNano:
    @pytest.fixture(params=["s", "ms", "us"])
    def unit(self, request):
        return request.param

    @pytest.fixture
    def reso(self, unit):
        if unit == "s":
            return NpyDatetimeUnit.NPY_FR_s.value
        elif unit == "ms":
            return NpyDatetimeUnit.NPY_FR_ms.value
        elif unit == "us":
            return NpyDatetimeUnit.NPY_FR_us.value
        else:
            raise NotImplementedError(unit)

    @pytest.fixture
    def tda(self, unit):
        arr = np.arange(5, dtype=np.int64).view(f"m8[{unit}]")
        return TimedeltaArray._simple_new(arr, dtype=arr.dtype)

    def test_non_nano(self, unit, reso):
        arr = np.arange(5, dtype=np.int64).view(f"m8[{unit}]")
        tda = TimedeltaArray._simple_new(arr, dtype=arr.dtype)

        assert tda.dtype == arr.dtype
        assert tda[0]._creso == reso

    @pytest.mark.parametrize("field", TimedeltaArray._field_ops)
    def test_fields(self, tda, field):
        as_nano = tda._ndarray.astype("m8[ns]")
        tda_nano = TimedeltaArray._simple_new(as_nano, dtype=as_nano.dtype)

        result = getattr(tda, field)
        expected = getattr(tda_nano, field)
        tm.assert_numpy_array_equal(result, expected)

    def test_to_pytimedelta(self, tda):
        as_nano = tda._ndarray.astype("m8[ns]")
        tda_nano = TimedeltaArray._simple_new(as_nano, dtype=as_nano.dtype)

        result = tda.to_pytimedelta()
        expected = tda_nano.to_pytimedelta()
        tm.assert_numpy_array_equal(result, expected)

    def test_total_seconds(self, unit, tda):
        as_nano = tda._ndarray.astype("m8[ns]")
        tda_nano = TimedeltaArray._simple_new(as_nano, dtype=as_nano.dtype)

        result = tda.total_seconds()
        expected = tda_nano.total_seconds()
        tm.assert_numpy_array_equal(result, expected)

    def test_timedelta_array_total_seconds(self):
        # GH34290
        expected = Timedelta("2 min").total_seconds()

        result = pd.array([Timedelta("2 min")]).total_seconds()[0]
        assert result == expected

    @pytest.mark.parametrize(
        "nat", [np.datetime64("NaT", "ns"), np.datetime64("NaT", "us")]
    )
    def test_add_nat_datetimelike_scalar(self, nat, tda):
        result = tda + nat
        assert isinstance(result, DatetimeArray)
        assert result._creso == tda._creso
        assert result.isna().all()

        result = nat + tda
        assert isinstance(result, DatetimeArray)
        assert result._creso == tda._creso
        assert result.isna().all()

    def test_add_pdnat(self, tda):
        result = tda + pd.NaT
        assert isinstance(result, TimedeltaArray)
        assert result._creso == tda._creso
        assert result.isna().all()

        result = pd.NaT + tda
        assert isinstance(result, TimedeltaArray)
        assert result._creso == tda._creso
        assert result.isna().all()

    # TODO: 2022-07-11 this is the only test that gets to DTA.tz_convert
    #  or tz_localize with non-nano; implement tests specific to that.
    def test_add_datetimelike_scalar(self, tda, tz_naive_fixture):
        ts = pd.Timestamp("2016-01-01", tz=tz_naive_fixture)

        expected = tda.as_unit("ns") + ts
        res = tda + ts
        tm.assert_extension_array_equal(res, expected)
        res = ts + tda
        tm.assert_extension_array_equal(res, expected)

        ts += Timedelta(1)  # case where we can't cast losslessly

        exp_values = tda._ndarray + ts.asm8
        expected = (
            DatetimeArray._simple_new(exp_values, dtype=exp_values.dtype)
            .tz_localize("UTC")
            .tz_convert(ts.tz)
        )

        result = tda + ts
        tm.assert_extension_array_equal(result, expected)

        result = ts + tda
        tm.assert_extension_array_equal(result, expected)

    def test_mul_scalar(self, tda):
        other = 2
        result = tda * other
        expected = TimedeltaArray._simple_new(tda._ndarray * other, dtype=tda.dtype)
        tm.assert_extension_array_equal(result, expected)
        assert result._creso == tda._creso

    def test_mul_listlike(self, tda):
        other = np.arange(len(tda))
        result = tda * other
        expected = TimedeltaArray._simple_new(tda._ndarray * other, dtype=tda.dtype)
        tm.assert_extension_array_equal(result, expected)
        assert result._creso == tda._creso

    def test_mul_listlike_object(self, tda):
        other = np.arange(len(tda))
        result = tda * other.astype(object)
        expected = TimedeltaArray._simple_new(tda._ndarray * other, dtype=tda.dtype)
        tm.assert_extension_array_equal(result, expected)
        assert result._creso == tda._creso

    def test_div_numeric_scalar(self, tda):
        other = 2
        result = tda / other
        expected = TimedeltaArray._simple_new(tda._ndarray / other, dtype=tda.dtype)
        tm.assert_extension_array_equal(result, expected)
        assert result._creso == tda._creso

    def test_div_td_scalar(self, tda):
        other = timedelta(seconds=1)
        result = tda / other
        expected = tda._ndarray / np.timedelta64(1, "s")
        tm.assert_numpy_array_equal(result, expected)

    def test_div_numeric_array(self, tda):
        other = np.arange(len(tda))
        result = tda / other
        expected = TimedeltaArray._simple_new(tda._ndarray / other, dtype=tda.dtype)
        tm.assert_extension_array_equal(result, expected)
        assert result._creso == tda._creso

    def test_div_td_array(self, tda):
        other = tda._ndarray + tda._ndarray[-1]
        result = tda / other
        expected = tda._ndarray / other
        tm.assert_numpy_array_equal(result, expected)

    def test_add_timedeltaarraylike(self, tda):
        tda_nano = tda.astype("m8[ns]")

        expected = tda_nano * 2
        res = tda_nano + tda
        tm.assert_extension_array_equal(res, expected)
        res = tda + tda_nano
        tm.assert_extension_array_equal(res, expected)

        expected = tda_nano * 0
        res = tda - tda_nano
        tm.assert_extension_array_equal(res, expected)

        res = tda_nano - tda
        tm.assert_extension_array_equal(res, expected)


class TestTimedeltaArray:
    @pytest.mark.parametrize("dtype", [int, np.int32, np.int64, "uint32", "uint64"])
    def test_astype_int(self, dtype):
        arr = TimedeltaArray._from_sequence([Timedelta("1H"), Timedelta("2H")])

        if np.dtype(dtype) != np.int64:
            with pytest.raises(TypeError, match=r"Do obj.astype\('int64'\)"):
                arr.astype(dtype)
            return

        result = arr.astype(dtype)
        expected = arr._ndarray.view("i8")
        tm.assert_numpy_array_equal(result, expected)

    def test_setitem_clears_freq(self):
        a = TimedeltaArray(pd.timedelta_range("1H", periods=2, freq="H"))
        a[0] = Timedelta("1H")
        assert a.freq is None

    @pytest.mark.parametrize(
        "obj",
        [
            Timedelta(seconds=1),
            Timedelta(seconds=1).to_timedelta64(),
            Timedelta(seconds=1).to_pytimedelta(),
        ],
    )
    def test_setitem_objects(self, obj):
        # make sure we accept timedelta64 and timedelta in addition to Timedelta
        tdi = pd.timedelta_range("2 Days", periods=4, freq="H")
        arr = TimedeltaArray(tdi, freq=tdi.freq)

        arr[0] = obj
        assert arr[0] == Timedelta(seconds=1)

    @pytest.mark.parametrize(
        "other",
        [
            1,
            np.int64(1),
            1.0,
            np.datetime64("NaT"),
            pd.Timestamp("2021-01-01"),
            "invalid",
            np.arange(10, dtype="i8") * 24 * 3600 * 10**9,
            (np.arange(10) * 24 * 3600 * 10**9).view("datetime64[ns]"),
            pd.Timestamp("2021-01-01").to_period("D"),
        ],
    )
    @pytest.mark.parametrize("index", [True, False])
    def test_searchsorted_invalid_types(self, other, index):
        data = np.arange(10, dtype="i8") * 24 * 3600 * 10**9
        arr = TimedeltaArray(data, freq="D")
        if index:
            arr = pd.Index(arr)

        msg = "|".join(
            [
                "searchsorted requires compatible dtype or scalar",
                "value should be a 'Timedelta', 'NaT', or array of those. Got",
            ]
        )
        with pytest.raises(TypeError, match=msg):
            arr.searchsorted(other)


class TestUnaryOps:
    def test_abs(self):
        vals = np.array([-3600 * 10**9, "NaT", 7200 * 10**9], dtype="m8[ns]")
        arr = TimedeltaArray(vals)

        evals = np.array([3600 * 10**9, "NaT", 7200 * 10**9], dtype="m8[ns]")
        expected = TimedeltaArray(evals)

        result = abs(arr)
        tm.assert_timedelta_array_equal(result, expected)

        result2 = np.abs(arr)
        tm.assert_timedelta_array_equal(result2, expected)

    def test_pos(self):
        vals = np.array([-3600 * 10**9, "NaT", 7200 * 10**9], dtype="m8[ns]")
        arr = TimedeltaArray(vals)

        result = +arr
        tm.assert_timedelta_array_equal(result, arr)
        assert not tm.shares_memory(result, arr)

        result2 = np.positive(arr)
        tm.assert_timedelta_array_equal(result2, arr)
        assert not tm.shares_memory(result2, arr)

    def test_neg(self):
        vals = np.array([-3600 * 10**9, "NaT", 7200 * 10**9], dtype="m8[ns]")
        arr = TimedeltaArray(vals)

        evals = np.array([3600 * 10**9, "NaT", -7200 * 10**9], dtype="m8[ns]")
        expected = TimedeltaArray(evals)

        result = -arr
        tm.assert_timedelta_array_equal(result, expected)

        result2 = np.negative(arr)
        tm.assert_timedelta_array_equal(result2, expected)

    def test_neg_freq(self):
        tdi = pd.timedelta_range("2 Days", periods=4, freq="H")
        arr = TimedeltaArray(tdi, freq=tdi.freq)

        expected = TimedeltaArray(-tdi._data, freq=-tdi.freq)

        result = -arr
        tm.assert_timedelta_array_equal(result, expected)

        result2 = np.negative(arr)
        tm.assert_timedelta_array_equal(result2, expected)