Allow passing datetime.date instances as timestamp input

Fixes: #288

Author: elprans

asyncpg/protocol/codecs/datetime.pyx

@@ -5,8 +5,11 @@
 # the Apache 2.0 License: http://www.apache.org/licenses/LICENSE-2.0
 
 
+cimport cpython.datetime
 import datetime
 
+cpython.datetime.import_datetime()
+
 utc = datetime.timezone.utc
 date_from_ordinal = datetime.date.fromordinal
 timedelta = datetime.timedelta
@@ -56,6 +59,11 @@ cdef int32_t negative_infinity_date_ord = <int32_t>cpython.PyLong_AsLong(
     negative_infinity_date.toordinal())
 
 
+cdef inline _local_timezone():
+    d = datetime.datetime.now(datetime.timezone.utc).astimezone()
+    return datetime.timezone(d.utcoffset())
+
+
 cdef inline _encode_time(WriteBuffer buf, int64_t seconds,
                          int32_t microseconds):
     # XXX: add support for double timestamps
@@ -135,6 +143,15 @@ cdef date_decode_tuple(ConnectionSettings settings, FastReadBuffer buf):
 
 
 cdef timestamp_encode(ConnectionSettings settings, WriteBuffer buf, obj):
+    if not cpython.datetime.PyDateTime_Check(obj):
+        if cpython.datetime.PyDate_Check(obj):
+            obj = datetime.datetime(obj.year, obj.month, obj.day)
+        else:
+            raise TypeError(
+                'expected a datetime.date or datetime.datetime instance, '
+                'got {!r}'.format(type(obj).__name__)
+            )
+
     delta = obj - pg_epoch_datetime
     cdef:
         int64_t seconds = cpython.PyLong_AsLongLong(delta.days) * 86400 + \
@@ -186,6 +203,16 @@ cdef timestamp_decode_tuple(ConnectionSettings settings, FastReadBuffer buf):
 
 
 cdef timestamptz_encode(ConnectionSettings settings, WriteBuffer buf, obj):
+    if not cpython.datetime.PyDateTime_Check(obj):
+        if cpython.datetime.PyDate_Check(obj):
+            obj = datetime.datetime(obj.year, obj.month, obj.day,
+                                    tzinfo=_local_timezone())
+        else:
+            raise TypeError(
+                'expected a datetime.date or datetime.datetime instance, '
+                'got {!r}'.format(type(obj).__name__)
+            )
+
     buf.write_int32(8)
 
     if obj == infinity_datetime:
@@ -195,7 +222,14 @@ cdef timestamptz_encode(ConnectionSettings settings, WriteBuffer buf, obj):
         buf.write_int64(pg_time64_negative_infinity)
         return
 
-    delta = obj.astimezone(utc) - pg_epoch_datetime_utc
+    try:
+        utc_dt = obj.astimezone(utc)
+    except ValueError:
+        # Python 3.5 doesn't like it when we call astimezone()
+        # on naive datetime objects, so make it aware.
+        utc_dt = obj.replace(tzinfo=_local_timezone()).astimezone(utc)
+
+    delta = utc_dt - pg_epoch_datetime_utc
     cdef:
         int64_t seconds = cpython.PyLong_AsLongLong(delta.days) * 86400 + \
                                 cpython.PyLong_AsLong(delta.seconds)

tests/test_codecs.py

@@ -25,13 +25,21 @@ def _timezone(offset):
     return datetime.timezone(datetime.timedelta(minutes=minutes))
 
 
+def _system_timezone():
+    d = datetime.datetime.now(datetime.timezone.utc).astimezone()
+    return datetime.timezone(d.utcoffset())
+
+
 infinity_datetime = datetime.datetime(
     datetime.MAXYEAR, 12, 31, 23, 59, 59, 999999)
 negative_infinity_datetime = datetime.datetime(
     datetime.MINYEAR, 1, 1, 0, 0, 0, 0)
 
 infinity_date = datetime.date(datetime.MAXYEAR, 12, 31)
 negative_infinity_date = datetime.date(datetime.MINYEAR, 1, 1)
+current_timezone = _system_timezone()
+current_date = datetime.date.today()
+current_datetime = datetime.datetime.now()
 
 
 type_samples = [
@@ -160,6 +168,8 @@ def _timezone(offset):
         negative_infinity_datetime,
         {'textinput': 'infinity', 'output': infinity_datetime},
         {'textinput': '-infinity', 'output': negative_infinity_datetime},
+        {'input': datetime.date(2000, 1, 1),
+         'output': datetime.datetime(2000, 1, 1)}
     ]),
     ('date', 'date', [
         datetime.date(3000, 5, 20),
@@ -185,6 +195,16 @@ def _timezone(offset):
         datetime.datetime(2400, 1, 1, 10, 10, 0, tzinfo=_timezone(2000)),
         infinity_datetime,
         negative_infinity_datetime,
+        {
+            'input': current_date,
+            'output': datetime.datetime(
+                year=current_date.year, month=current_date.month,
+                day=current_date.day, tzinfo=current_timezone),
+        },
+        {
+            'input': current_datetime,
+            'output': current_datetime.replace(tzinfo=current_timezone),
+        }
     ]),
     ('timetz', 'timetz', [
         # timetz retains the offset
@@ -657,6 +677,12 @@ async def test_invalid_input(self):
                 2 ** 32,
                 -1,
             ]),
+            ('timestamp', r"expected a datetime\.date.*got 'str'", [
+                'foo'
+            ]),
+            ('timestamptz', r"expected a datetime\.date.*got 'str'", [
+                'foo'
+            ]),
         ]
 
         for typname, errmsg, data in cases: