diff --git a/mp_yearmonth/yearmonth.py b/mp_yearmonth/yearmonth.py index e4c9bd5..850f160 100644 --- a/mp_yearmonth/yearmonth.py +++ b/mp_yearmonth/yearmonth.py @@ -108,6 +108,11 @@ def __sub__(self, other) -> "YearMonth": return self.applying_delta(-other) return NotImplemented + def __contains__(self, other) -> bool: + if isinstance(other, datetime.date) or isinstance(other, datetime.datetime): + return self.year == other.year and self.month == other.month + return False + # Public interface @property diff --git a/tests/test_yearmonth.py b/tests/test_yearmonth.py index 3fbc2f0..c372575 100644 --- a/tests/test_yearmonth.py +++ b/tests/test_yearmonth.py @@ -185,6 +185,24 @@ def test_sub(): assert ym - 2 == YearMonth(2020, 11) +def test_contains_handles_dates(): + assert datetime.date(2021, 1, 1) in YearMonth(2021, 1) + assert datetime.date(2021, 1, 31) in YearMonth(2021, 1) + assert datetime.date(2021, 2, 1) not in YearMonth(2021, 1) + + +def test_contains_handles_datetimes(): + assert datetime.datetime(2021, 1, 1, hour=12, minute=30) in YearMonth(2021, 1) + assert datetime.datetime(2021, 1, 31, hour=12, minute=30) in YearMonth(2021, 1) + assert datetime.datetime(2021, 2, 1, hour=12, minute=30) not in YearMonth(2021, 1) + + +def test_contains_handles_unsupported_types(): + assert 1 not in YearMonth(2021, 1) + assert 2021 not in YearMonth(2021, 1) + assert "test" not in YearMonth(2021, 1) + + def test_bounds(): ym1 = YearMonth(2021, 1) assert ym1.bounds() == (datetime.date(2021, 1, 1), datetime.date(2021, 1, 31))