Introduce Date type to accomodate ranges outside datetime.[MIN|MAX}YEAR

aholmberg · aholmberg · commit 4f3c77c970e1 · 2015-03-13T15:02:31.000-05:00
diff --git a/cassandra/cqltypes.py b/cassandra/cqltypes.py
@@ -31,7 +31,6 @@
 from binascii import unhexlify
 import calendar
 from collections import namedtuple
-import datetime
 from decimal import Decimal
 import io
 import logging
@@ -626,38 +625,29 @@ def serialize(timeuuid, protocol_version):
 
 class SimpleDateType(_CassandraType):
     typename = 'date'
-    seconds_per_day = 60 * 60 * 24
     date_format = "%Y-%m-%d"
 
     @classmethod
     def validate(cls, val):
-        if isinstance(val, six.string_types):
-            val = cls.interpret_simpledate_string(val)
-        elif (not isinstance(val, datetime.date)) and not isinstance(val, six.integer_types):
-            raise TypeError('SimpleDateType arg must be a datetime.date, unsigned integer, or string in the format YYYY-MM-DD')
+        if not isinstance(val, util.Date):
+            val = util.Date(val)
         return val
 
-    @staticmethod
-    def interpret_simpledate_string(v):
-        date_time = datetime.datetime.strptime(v, SimpleDateType.date_format)
-        return datetime.date(date_time.year, date_time.month, date_time.day)
-
     @staticmethod
     def serialize(val, protocol_version):
         # Values of the 'date'` type are encoded as 32-bit unsigned integers
-        # representing a number of days with "the epoch" at the center of the
-        # range (2^31). Epoch is January 1st, 1970
+        # representing a number of days with epoch (January 1st, 1970) at the center of the
+        # range (2^31).
         try:
-            shifted = (calendar.timegm(val.timetuple()) // SimpleDateType.seconds_per_day) + 2 ** 31
+            days = val.days_from_epoch
         except AttributeError:
-            shifted = val
-        return uint32_pack(shifted)
+            days = util.Date(val).days_from_epoch
+        return uint32_pack(days + 2 ** 31)
 
     @staticmethod
     def deserialize(byts, protocol_version):
-        timestamp = SimpleDateType.seconds_per_day * (uint32_unpack(byts) - 2 ** 31)
-        dt = util.datetime_from_timestamp(timestamp)
-        return datetime.date(dt.year, dt.month, dt.day)
+        days = uint32_unpack(byts) - 2 ** 31
+        return util.Date(days)
 
 
 class TimeType(_CassandraType):
diff --git a/cassandra/util.py b/cassandra/util.py
@@ -892,3 +892,74 @@ def __repr__(self):
     def __str__(self):
         return "%02d:%02d:%02d.%09d" % (self.hour, self.minute,
                                         self.second, self.nanosecond)
+
+
+class Date(object):
+    '''
+    Idealized naive date: year, month, day
+
+    Offers wider year range than datetime.date. Dates that cannot be represented
+    as a date (because datetime.MINYEAR, datetime.MAXYEAR), this type falls back
+    to printing days_from_epoch offset.
+    '''
+
+    MINUTE = 60
+    HOUR = 60 * MINUTE
+    DAY = 24 * HOUR
+
+    date_format = "%Y-%m-%d"
+
+    days_from_epoch = 0
+
+    def __init__(self, value):
+        if isinstance(value, six.integer_types):
+            self.days_from_epoch = value
+        elif isinstance(value, (datetime.date, datetime.datetime)):
+            self._from_timetuple(value.timetuple())
+        elif isinstance(value, six.string_types):
+            self._from_datestring(value)
+        else:
+            raise TypeError('Date arguments must be a whole number, datetime.date, or string')
+
+    @property
+    def seconds(self):
+        return self.days_from_epoch * Date.DAY
+
+    def date(self):
+        try:
+            dt = datetime_from_timestamp(self.seconds)
+            return datetime.date(dt.year, dt.month, dt.day)
+        except Exception:
+            raise ValueError("%r exceeds ranges for built-in datetime.date" % self)
+
+    def _from_timetuple(self, t):
+        self.days_from_epoch = calendar.timegm(t) // Date.DAY
+
+    def _from_datestring(self, s):
+        if s[0] == '+':
+            s = s[1:]
+        dt = datetime.datetime.strptime(s, self.date_format)
+        self._from_timetuple(dt.timetuple())
+
+    def __eq__(self, other):
+        if isinstance(other, Date):
+            return self.days_from_epoch == other.days_from_epoch
+
+        if isinstance(other, six.integer_types):
+            return self.days_from_epoch == other
+
+        try:
+            return self.date == other
+        except Exception:
+            return False
+
+    def __repr__(self):
+        return "Date(%s)" % self.days_from_epoch
+
+    def __str__(self):
+        try:
+            dt = datetime_from_timestamp(self.seconds)
+            return "%04d-%02d-%02d" % (dt.year, dt.month, dt.day)
+        except:
+            # If we overflow datetime.[MIN|M
+            return str(self.days_from_epoch)
diff --git a/tests/unit/test_marshalling.py b/tests/unit/test_marshalling.py
@@ -24,7 +24,7 @@
 from uuid import UUID
 
 from cassandra.cqltypes import lookup_casstype, DecimalType, UTF8Type
-from cassandra.util import OrderedMap, OrderedMapSerializedKey, sortedset, Time
+from cassandra.util import OrderedMap, OrderedMapSerializedKey, sortedset, Time, Date
 
 marshalled_value_pairs = (
     # binary form, type, python native type
@@ -79,8 +79,8 @@
     (b'\x00\x00', 'ListType(FloatType)', []),
     (b'\x00\x00', 'SetType(IntegerType)', sortedset()),
     (b'\x00\x01\x00\x10\xafYC\xa3\xea<\x11\xe1\xabc\xc4,\x03"y\xf0', 'ListType(TimeUUIDType)', [UUID(bytes=b'\xafYC\xa3\xea<\x11\xe1\xabc\xc4,\x03"y\xf0')]),
-    (b'\x80\x00\x00\x01', 'SimpleDateType', date(1970,1,2)),
-    (b'\x7f\xff\xff\xff', 'SimpleDateType', date(1969,12,31)),
+    (b'\x80\x00\x00\x01', 'SimpleDateType', Date(1)),
+    (b'\x7f\xff\xff\xff', 'SimpleDateType', Date('1969-12-31')),
     (b'\x00\x00\x00\x00\x00\x00\x00\x01', 'TimeType', Time(1))
 )
 
diff --git a/tests/unit/test_types.py b/tests/unit/test_types.py
@@ -130,7 +130,7 @@ def test_simpledate(self):
         Test cassandra.cqltypes.SimpleDateType() construction
         """
         # from string
-        expected_date = datetime.date(1492, 10, 12)
+        expected_date = cassandra.util.Date(datetime.date(1492, 10, 12))
         sd = SimpleDateType('1492-10-12')
         self.assertEqual(sd.val, expected_date)
 
@@ -139,7 +139,7 @@ def test_simpledate(self):
         self.assertEqual(sd.val, expected_date)
 
         # int
-        expected_timestamp = calendar.timegm(expected_date.timetuple())
+        expected_timestamp = calendar.timegm(expected_date.date().timetuple())
         sd = SimpleDateType(expected_timestamp)
         self.assertEqual(sd.val, expected_timestamp)
 
