diff --git a/src/undate/date.py b/src/undate/date.py index 4422a71..c6a095f 100644 --- a/src/undate/date.py +++ b/src/undate/date.py @@ -17,11 +17,11 @@ class Date(np.ndarray): # extend np.datetime64 datatype # adapted from https://stackoverflow.com/a/27129510/9706217 - def __new__(cls, year: str, month: str = None, day: str = None): + def __new__(cls, year: int, month: int = None, day: int = None): if isinstance(year, np.datetime64): data = year else: - datestr = year + datestr = str(year) if month is not None: datestr = f"{year}-{month:02d}" if day is not None: @@ -29,9 +29,18 @@ def __new__(cls, year: str, month: str = None, day: str = None): data = np.datetime64(datestr) data = np.asarray(data, dtype="datetime64") - if data.dtype != "datetime64[D]": + + # expected format depends on granularity / how much of date is known + expected_granularity = "Y" + if day is not None and month is not None: + expected_granularity = "D" + elif month: + expected_granularity = "M" + expected_dtype = f"datetime64[{expected_granularity}]" + + if data.dtype != expected_dtype: raise Exception( - "Unable to parse dates adequately to datetime64[D]: %s" % data + f"Unable to parse dates adequately as {expected_dtype}: {data}" ) obj = data.view(cls) return obj @@ -51,11 +60,15 @@ def year(self): @property def month(self): - return int(str(self.astype("datetime64[M]")).split("-")[-1]) + # if date unit is year, don't return a month (only M/D) + if not self.dtype == "datetime64[Y]": + return int(str(self.astype("datetime64[M]")).split("-")[-1]) @property def day(self): - return int(str(self.astype("datetime64[D]")).split("-")[-1]) + # only return a day if date unit is in days + if self.dtype == "datetime64[D]": + return int(str(self.astype("datetime64[D]")).split("-")[-1]) class DatePrecision(IntEnum): @@ -75,5 +88,5 @@ class DatePrecision(IntEnum): def __str__(self): return f"{self.name}" - # NOTE: consider harmonizing / using numpy units - # numpy date units are years (‘Y’), months (‘M’), weeks (‘W’), and days (‘D’) + # NOTE: consider harmonizing / using numpy date units: + # years (‘Y’), months (‘M’), weeks (‘W’), and days (‘D’) diff --git a/tests/test_date.py b/tests/test_date.py new file mode 100644 index 0000000..cb56790 --- /dev/null +++ b/tests/test_date.py @@ -0,0 +1,43 @@ +import numpy as np + +from undate.date import Date, DatePrecision, ONE_DAY, ONE_YEAR, ONE_MONTH_MAX + + +class TestDatePrecision: + def test_str(self): + assert str(DatePrecision.YEAR) == "YEAR" + + +class TestDate: + def test_init_year(self): + d = Date(2001) + assert isinstance(d, Date) + assert d.dtype == "datetime64[Y]" + assert str(d) == "2001" + + def test_init_year_month(self): + d = Date(2010, 5) + assert isinstance(d, Date) + assert d.dtype == "datetime64[M]" + assert str(d) == "2010-05" + + def test_init_year_month(self): + d = Date(2021, 6, 15) + assert isinstance(d, Date) + assert d.dtype == "datetime64[D]" + assert str(d) == "2021-06-15" + + def test_properties_year(self): + assert Date(2001).year == 2001 + assert Date(2010, 5).year == 2010 + assert Date(2021, 6, 15).year == 2021 + + def test_properties_month(self): + assert Date(2001).month is None + assert Date(2010, 5).month == 5 + assert Date(2021, 6, 15).month == 6 + + def test_properties_day(self): + assert Date(2001).day is None + assert Date(2010, 5).day == None + assert Date(2021, 6, 15).day == 15 diff --git a/tests/test_undate.py b/tests/test_undate.py index 7c4c577..e2ab201 100644 --- a/tests/test_undate.py +++ b/tests/test_undate.py @@ -2,12 +2,7 @@ import pytest -from undate.undate import Undate, UndateInterval, DatePrecision - - -class TestDatePrecision: - def test_str(self): - assert str(DatePrecision.YEAR) == "YEAR" +from undate.undate import Undate, UndateInterval class TestUndate: