Add support for membership operators (#4)

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

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))