Add unit tests for date class

src/undate/date.py

@@ -17,21 +17,30 @@ 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:
                     datestr = f"{datestr}-{day:02d}"
             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’)

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

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: