diff --git a/src/hamster/lib/datetime.py b/src/hamster/lib/datetime.py index cb4d84618..ce9a303b6 100644 --- a/src/hamster/lib/datetime.py +++ b/src/hamster/lib/datetime.py @@ -46,6 +46,21 @@ class date(pdt.date): def __new__(cls, year, month, day): return pdt.date.__new__(cls, year, month, day) + def __add__(self, other): + # python date.__add__ was not type stable prior to 3.8 + return self.from_pdt(self.to_pdt() + other) + + __radd__ = __add__ + + def __sub__(self, other): + # python date.__sub__ was not type stable prior to 3.8 + if isinstance(other, timedelta): + return self.from_pdt(self.to_pdt() - other) + elif isinstance(other, date): + return timedelta.from_pdt(self.to_pdt() - other) + else: + raise NotImplementedError("subtract {}".format(type(other))) + @classmethod def parse(cls, s): """Return date from string.""" @@ -73,6 +88,10 @@ def from_pdt(cls, d): """Convert python date to hamster date.""" return cls(d.year, d.month, d.day) + def to_pdt(self): + """Convert to python date.""" + return pdt.date(self.year, self.month, self.day) + # For datetime that will need to be outside the class. # Same here for consistency date.re = re.compile(date.pattern(), flags=re.VERBOSE) diff --git a/tests/stuff_test.py b/tests/stuff_test.py index e22e3dae9..c0b3088cc 100644 --- a/tests/stuff_test.py +++ b/tests/stuff_test.py @@ -281,6 +281,8 @@ def test_datetime_hday(self): date_time = dt.datetime(2018, 8, 14, 0, 10) # 2018-08-14 0:10 expected = dt.date(2018, 8, 13) self.assertEqual(date_time.hday(), expected) + today = dt.hday.today() + self.assertEqual(type(today), dt.hday) def test_parse_date(self): date = dt.date.parse("2020-01-05")