keshe/.venv/Lib/site-packages/pandas/tests/scalar/interval/test_arithmetic.py

from datetime import timedelta

import numpy as np
import pytest

from pandas import (
    Interval,
    Timedelta,
    Timestamp,
)
import pandas._testing as tm


class TestIntervalArithmetic:
    def test_interval_add(self, closed):
        interval = Interval(0, 1, closed=closed)
        expected = Interval(1, 2, closed=closed)

        result = interval + 1
        assert result == expected

        result = 1 + interval
        assert result == expected

        result = interval
        result += 1
        assert result == expected

        msg = r"unsupported operand type\(s\) for \+"
        with pytest.raises(TypeError, match=msg):
            interval + interval

        with pytest.raises(TypeError, match=msg):
            interval + "foo"

    def test_interval_sub(self, closed):
        interval = Interval(0, 1, closed=closed)
        expected = Interval(-1, 0, closed=closed)

        result = interval - 1
        assert result == expected

        result = interval
        result -= 1
        assert result == expected

        msg = r"unsupported operand type\(s\) for -"
        with pytest.raises(TypeError, match=msg):
            interval - interval

        with pytest.raises(TypeError, match=msg):
            interval - "foo"

    def test_interval_mult(self, closed):
        interval = Interval(0, 1, closed=closed)
        expected = Interval(0, 2, closed=closed)

        result = interval * 2
        assert result == expected

        result = 2 * interval
        assert result == expected

        result = interval
        result *= 2
        assert result == expected

        msg = r"unsupported operand type\(s\) for \*"
        with pytest.raises(TypeError, match=msg):
            interval * interval

        msg = r"can\'t multiply sequence by non-int"
        with pytest.raises(TypeError, match=msg):
            interval * "foo"

    def test_interval_div(self, closed):
        interval = Interval(0, 1, closed=closed)
        expected = Interval(0, 0.5, closed=closed)

        result = interval / 2.0
        assert result == expected

        result = interval
        result /= 2.0
        assert result == expected

        msg = r"unsupported operand type\(s\) for /"
        with pytest.raises(TypeError, match=msg):
            interval / interval

        with pytest.raises(TypeError, match=msg):
            interval / "foo"

    def test_interval_floordiv(self, closed):
        interval = Interval(1, 2, closed=closed)
        expected = Interval(0, 1, closed=closed)

        result = interval // 2
        assert result == expected

        result = interval
        result //= 2
        assert result == expected

        msg = r"unsupported operand type\(s\) for //"
        with pytest.raises(TypeError, match=msg):
            interval // interval

        with pytest.raises(TypeError, match=msg):
            interval // "foo"

    @pytest.mark.parametrize("method", ["__add__", "__sub__"])
    @pytest.mark.parametrize(
        "interval",
        [
            Interval(
                Timestamp("2017-01-01 00:00:00"), Timestamp("2018-01-01 00:00:00")
            ),
            Interval(Timedelta(days=7), Timedelta(days=14)),
        ],
    )
    @pytest.mark.parametrize(
        "delta", [Timedelta(days=7), timedelta(7), np.timedelta64(7, "D")]
    )
    def test_time_interval_add_subtract_timedelta(self, interval, delta, method):
        # https://github.com/pandas-dev/pandas/issues/32023
        result = getattr(interval, method)(delta)
        left = getattr(interval.left, method)(delta)
        right = getattr(interval.right, method)(delta)
        expected = Interval(left, right)

        assert result == expected

    @pytest.mark.parametrize("interval", [Interval(1, 2), Interval(1.0, 2.0)])
    @pytest.mark.parametrize(
        "delta", [Timedelta(days=7), timedelta(7), np.timedelta64(7, "D")]
    )
    def test_numeric_interval_add_timedelta_raises(self, interval, delta):
        # https://github.com/pandas-dev/pandas/issues/32023
        msg = "|".join(
            [
                "unsupported operand",
                "cannot use operands",
                "Only numeric, Timestamp and Timedelta endpoints are allowed",
            ]
        )
        with pytest.raises((TypeError, ValueError), match=msg):
            interval + delta

        with pytest.raises((TypeError, ValueError), match=msg):
            delta + interval

    @pytest.mark.parametrize("klass", [timedelta, np.timedelta64, Timedelta])
    def test_timedelta_add_timestamp_interval(self, klass):
        delta = klass(0)
        expected = Interval(Timestamp("2020-01-01"), Timestamp("2020-02-01"))

        result = delta + expected
        assert result == expected

        result = expected + delta
        assert result == expected


class TestIntervalComparisons:
    def test_interval_equal(self):
        assert Interval(0, 1) == Interval(0, 1, closed="right")
        assert Interval(0, 1) != Interval(0, 1, closed="left")
        assert Interval(0, 1) != 0

    def test_interval_comparison(self):
        msg = (
            "'<' not supported between instances of "
            "'pandas._libs.interval.Interval' and 'int'"
        )
        with pytest.raises(TypeError, match=msg):
            Interval(0, 1) < 2

        assert Interval(0, 1) < Interval(1, 2)
        assert Interval(0, 1) < Interval(0, 2)
        assert Interval(0, 1) < Interval(0.5, 1.5)
        assert Interval(0, 1) <= Interval(0, 1)
        assert Interval(0, 1) > Interval(-1, 2)
        assert Interval(0, 1) >= Interval(0, 1)

    def test_equality_comparison_broadcasts_over_array(self):
        # https://github.com/pandas-dev/pandas/issues/35931
        interval = Interval(0, 1)
        arr = np.array([interval, interval])
        result = interval == arr
        expected = np.array([True, True])
        tm.assert_numpy_array_equal(result, expected)