Return objects with timezone parsing infinity timestamptz

Close #536.

6 files changed, 160 insertions, 67 deletions

diff --git a/NEWS b/NEWS
index b2d2153..84be201 100644
--- a/NEWS
+++ b/NEWS
@@ -9,6 +9,8 @@ What's new in psycopg 2.7.2
   2.7 by mistake.
 - Don't display the password in `connection.dsn` when the connection
   string is specified as an URI (:ticket:`#528`).
+- Return objects with timezone parsing "infinity" :sql:`timestamptz`
+  (:ticket:`#536`).
 
 
 What's new in psycopg 2.7.1
diff --git a/psycopg/typecast.c b/psycopg/typecast.c
index 4fd92be..cf19fcd 100644
--- a/psycopg/typecast.c
+++ b/psycopg/typecast.c
@@ -197,6 +197,7 @@ typecast_UNKNOWN_cast(const char *str, Py_ssize_t len, PyObject *curs)
 #include "psycopg/typecast_builtins.c"
 
 #define typecast_PYDATETIMEARRAY_cast typecast_GENERIC_ARRAY_cast
+#define typecast_PYDATETIMETZARRAY_cast typecast_GENERIC_ARRAY_cast
 #define typecast_PYDATEARRAY_cast typecast_GENERIC_ARRAY_cast
 #define typecast_PYTIMEARRAY_cast typecast_GENERIC_ARRAY_cast
 #define typecast_PYINTERVALARRAY_cast typecast_GENERIC_ARRAY_cast
@@ -204,10 +205,12 @@ typecast_UNKNOWN_cast(const char *str, Py_ssize_t len, PyObject *curs)
 /* a list of initializers, used to make the typecasters accessible anyway */
 static typecastObject_initlist typecast_pydatetime[] = {
     {"PYDATETIME", typecast_DATETIME_types, typecast_PYDATETIME_cast},
+    {"PYDATETIMETZ", typecast_DATETIMETZ_types, typecast_PYDATETIMETZ_cast},
     {"PYTIME", typecast_TIME_types, typecast_PYTIME_cast},
     {"PYDATE", typecast_DATE_types, typecast_PYDATE_cast},
     {"PYINTERVAL", typecast_INTERVAL_types, typecast_PYINTERVAL_cast},
     {"PYDATETIMEARRAY", typecast_DATETIMEARRAY_types, typecast_PYDATETIMEARRAY_cast, "PYDATETIME"},
+    {"PYDATETIMETZARRAY", typecast_DATETIMETZARRAY_types, typecast_PYDATETIMETZARRAY_cast, "PYDATETIMETZ"},
     {"PYTIMEARRAY", typecast_TIMEARRAY_types, typecast_PYTIMEARRAY_cast, "PYTIME"},
     {"PYDATEARRAY", typecast_DATEARRAY_types, typecast_PYDATEARRAY_cast, "PYDATE"},
     {"PYINTERVALARRAY", typecast_INTERVALARRAY_types, typecast_PYINTERVALARRAY_cast, "PYINTERVAL"},
@@ -216,6 +219,7 @@ static typecastObject_initlist typecast_pydatetime[] = {
 
 #ifdef HAVE_MXDATETIME
 #define typecast_MXDATETIMEARRAY_cast typecast_GENERIC_ARRAY_cast
+#define typecast_MXDATETIMETZARRAY_cast typecast_GENERIC_ARRAY_cast
 #define typecast_MXDATEARRAY_cast typecast_GENERIC_ARRAY_cast
 #define typecast_MXTIMEARRAY_cast typecast_GENERIC_ARRAY_cast
 #define typecast_MXINTERVALARRAY_cast typecast_GENERIC_ARRAY_cast
@@ -223,10 +227,12 @@ static typecastObject_initlist typecast_pydatetime[] = {
 /* a list of initializers, used to make the typecasters accessible anyway */
 static typecastObject_initlist typecast_mxdatetime[] = {
     {"MXDATETIME", typecast_DATETIME_types, typecast_MXDATE_cast},
+    {"MXDATETIMETZ", typecast_DATETIMETZ_types, typecast_MXDATE_cast},
     {"MXTIME", typecast_TIME_types, typecast_MXTIME_cast},
     {"MXDATE", typecast_DATE_types, typecast_MXDATE_cast},
     {"MXINTERVAL", typecast_INTERVAL_types, typecast_MXINTERVAL_cast},
     {"MXDATETIMEARRAY", typecast_DATETIMEARRAY_types, typecast_MXDATETIMEARRAY_cast, "MXDATETIME"},
+    {"MXDATETIMETZARRAY", typecast_DATETIMETZARRAY_types, typecast_MXDATETIMETZARRAY_cast, "MXDATETIMETZ"},
     {"MXTIMEARRAY", typecast_TIMEARRAY_types, typecast_MXTIMEARRAY_cast, "MXTIME"},
     {"MXDATEARRAY", typecast_DATEARRAY_types, typecast_MXDATEARRAY_cast, "MXDATE"},
     {"MXINTERVALARRAY", typecast_INTERVALARRAY_types, typecast_MXINTERVALARRAY_cast, "MXINTERVAL"},
diff --git a/psycopg/typecast_builtins.c b/psycopg/typecast_builtins.c
index fa548a7..75bcaec 100644
--- a/psycopg/typecast_builtins.c
+++ b/psycopg/typecast_builtins.c
@@ -7,6 +7,7 @@ static long int typecast_UNICODE_types[] = {19, 18, 25, 1042, 1043, 0};
 static long int typecast_STRING_types[] = {19, 18, 25, 1042, 1043, 0};
 static long int typecast_BOOLEAN_types[] = {16, 0};
 static long int typecast_DATETIME_types[] = {1114, 1184, 704, 1186, 0};
+static long int typecast_DATETIMETZ_types[] = {1184, 0};
 static long int typecast_TIME_types[] = {1083, 1266, 0};
 static long int typecast_DATE_types[] = {1082, 0};
 static long int typecast_INTERVAL_types[] = {704, 1186, 0};
@@ -20,6 +21,7 @@ static long int typecast_UNICODEARRAY_types[] = {1002, 1003, 1009, 1014, 1015, 0
 static long int typecast_STRINGARRAY_types[] = {1002, 1003, 1009, 1014, 1015, 0};
 static long int typecast_BOOLEANARRAY_types[] = {1000, 0};
 static long int typecast_DATETIMEARRAY_types[] = {1115, 1185, 0};
+static long int typecast_DATETIMETZARRAY_types[] = {1185, 0};
 static long int typecast_TIMEARRAY_types[] = {1183, 1270, 0};
 static long int typecast_DATEARRAY_types[] = {1182, 0};
 static long int typecast_INTERVALARRAY_types[] = {1187, 0};
@@ -41,6 +43,7 @@ static typecastObject_initlist typecast_builtins[] = {
   {"STRING", typecast_STRING_types, typecast_STRING_cast, NULL},
   {"BOOLEAN", typecast_BOOLEAN_types, typecast_BOOLEAN_cast, NULL},
   {"DATETIME", typecast_DATETIME_types, typecast_DATETIME_cast, NULL},
+  {"DATETIMETZ", typecast_DATETIMETZ_types, typecast_DATETIMETZ_cast, NULL},
   {"TIME", typecast_TIME_types, typecast_TIME_cast, NULL},
   {"DATE", typecast_DATE_types, typecast_DATE_cast, NULL},
   {"INTERVAL", typecast_INTERVAL_types, typecast_INTERVAL_cast, NULL},
diff --git a/psycopg/typecast_datetime.c b/psycopg/typecast_datetime.c
index a833d86..fec57de 100644
--- a/psycopg/typecast_datetime.c
+++ b/psycopg/typecast_datetime.c
@@ -80,91 +80,152 @@ typecast_PYDATE_cast(const char *str, Py_ssize_t len, PyObject *curs)
     return obj;
 }
 
-/** DATETIME - cast a timestamp into a datetime python object **/
+/* convert the strings -infinity and infinity into a datetime with timezone */
+static PyObject *
+_parse_inftz(const char *str, PyObject *curs)
+{
+    PyObject *rv = NULL;
+    PyObject *m = NULL;
+    PyObject *tzinfo_factory = NULL;
+    PyObject *tzinfo = NULL;
+    PyObject *args = NULL;
+    PyObject *kwargs = NULL;
+    PyObject *replace = NULL;
+
+    if (!(m = PyObject_GetAttrString(
+            (PyObject*)PyDateTimeAPI->DateTimeType,
+            (str[0] == '-' ? "min" : "max")))) {
+        goto exit;
+    }
+
+    tzinfo_factory = ((cursorObject *)curs)->tzinfo_factory;
+    if (tzinfo_factory == Py_None) {
+        rv = m;
+        m = NULL;
+        goto exit;
+    }
+
+    if (!(tzinfo = PyObject_CallFunction(tzinfo_factory, "i", 0))) {
+        goto exit;
+    }
+
+    /* m.replace(tzinfo=tzinfo) */
+    if (!(args = PyTuple_New(0))) { goto exit; }
+    if (!(kwargs = PyDict_New())) { goto exit; }
+    if (0 != PyDict_SetItemString(kwargs, "tzinfo", tzinfo)) { goto exit; }
+    if (!(replace = PyObject_GetAttrString(m, "replace"))) { goto exit; }
+    rv = PyObject_Call(replace, args, kwargs);
+
+exit:
+    Py_XDECREF(replace);
+    Py_XDECREF(args);
+    Py_XDECREF(kwargs);
+    Py_XDECREF(tzinfo);
+    Py_XDECREF(m);
+
+    return rv;
+}
 
 static PyObject *
-typecast_PYDATETIME_cast(const char *str, Py_ssize_t len, PyObject *curs)
+_parse_noninftz(const char *str, Py_ssize_t len, PyObject *curs)
 {
-    PyObject* obj = NULL;
+    PyObject* rv = NULL;
     PyObject *tzinfo = NULL;
     PyObject *tzinfo_factory;
     int n, y=0, m=0, d=0;
     int hh=0, mm=0, ss=0, us=0, tz=0;
     const char *tp = NULL;
 
+    Dprintf("typecast_PYDATETIMETZ_cast: s = %s", str);
+    n = typecast_parse_date(str, &tp, &len, &y, &m, &d);
+    Dprintf("typecast_PYDATE_cast: tp = %p "
+            "n = %d, len = " FORMAT_CODE_PY_SSIZE_T ","
+            " y = %d, m = %d, d = %d",
+             tp, n, len, y, m, d);
+    if (n != 3) {
+        PyErr_SetString(DataError, "unable to parse date");
+        goto exit;
+    }
+
+    if (len > 0) {
+        n = typecast_parse_time(tp, NULL, &len, &hh, &mm, &ss, &us, &tz);
+        Dprintf("typecast_PYDATETIMETZ_cast: n = %d,"
+            " len = " FORMAT_CODE_PY_SSIZE_T ","
+            " hh = %d, mm = %d, ss = %d, us = %d, tz = %d",
+            n, len, hh, mm, ss, us, tz);
+        if (n < 3 || n > 6) {
+            PyErr_SetString(DataError, "unable to parse time");
+            goto exit;
+        }
+    }
+
+    if (ss > 59) {
+        mm += 1;
+        ss -= 60;
+    }
+    if (y > 9999)
+        y = 9999;
+
+    tzinfo_factory = ((cursorObject *)curs)->tzinfo_factory;
+    if (n >= 5 && tzinfo_factory != Py_None) {
+        /* we have a time zone, calculate minutes and create
+           appropriate tzinfo object calling the factory */
+        Dprintf("typecast_PYDATETIMETZ_cast: UTC offset = %ds", tz);
+
+        /* The datetime module requires that time zone offsets be
+           a whole number of minutes, so truncate the seconds to the
+           closest minute. */
+        // printf("%d %d %d\n", tz, tzmin, round(tz / 60.0));
+        if (!(tzinfo = PyObject_CallFunction(tzinfo_factory, "i",
+                (int)round(tz / 60.0)))) {
+            goto exit;
+        }
+    } else {
+        Py_INCREF(Py_None);
+        tzinfo = Py_None;
+    }
+
+    Dprintf("typecast_PYDATETIMETZ_cast: tzinfo: %p, refcnt = "
+        FORMAT_CODE_PY_SSIZE_T,
+        tzinfo, Py_REFCNT(tzinfo));
+    rv = PyObject_CallFunction(
+        (PyObject*)PyDateTimeAPI->DateTimeType, "iiiiiiiO",
+        y, m, d, hh, mm, ss, us, tzinfo);
+
+exit:
+    Py_XDECREF(tzinfo);
+    return rv;
+}
+
+/** DATETIME - cast a timestamp into a datetime python object **/
+
+static PyObject *
+typecast_PYDATETIME_cast(const char *str, Py_ssize_t len, PyObject *curs)
+{
     if (str == NULL) { Py_RETURN_NONE; }
 
     /* check for infinity */
     if (!strcmp(str, "infinity") || !strcmp(str, "-infinity")) {
-        if (str[0] == '-') {
-            obj = PyObject_GetAttrString(
-                (PyObject*)PyDateTimeAPI->DateTimeType, "min");
-        }
-        else {
-            obj = PyObject_GetAttrString(
-                (PyObject*)PyDateTimeAPI->DateTimeType, "max");
-        }
+        return PyObject_GetAttrString(
+            (PyObject*)PyDateTimeAPI->DateTimeType,
+            (str[0] == '-' ? "min" : "max"));
     }
 
-    else {
-        Dprintf("typecast_PYDATETIME_cast: s = %s", str);
-        n = typecast_parse_date(str, &tp, &len, &y, &m, &d);
-        Dprintf("typecast_PYDATE_cast: tp = %p "
-                "n = %d, len = " FORMAT_CODE_PY_SSIZE_T ","
-                " y = %d, m = %d, d = %d",
-                 tp, n, len, y, m, d);
-        if (n != 3) {
-            PyErr_SetString(DataError, "unable to parse date");
-            return NULL;
-        }
+    return _parse_noninftz(str, len, curs);
+}
 
-        if (len > 0) {
-            n = typecast_parse_time(tp, NULL, &len, &hh, &mm, &ss, &us, &tz);
-            Dprintf("typecast_PYDATETIME_cast: n = %d,"
-                " len = " FORMAT_CODE_PY_SSIZE_T ","
-                " hh = %d, mm = %d, ss = %d, us = %d, tz = %d",
-                n, len, hh, mm, ss, us, tz);
-            if (n < 3 || n > 6) {
-                PyErr_SetString(DataError, "unable to parse time");
-                return NULL;
-            }
-        }
+/** DATETIMETZ - cast a timestamptz into a datetime python object **/
 
-        if (ss > 59) {
-            mm += 1;
-            ss -= 60;
-        }
-        if (y > 9999)
-            y = 9999;
-
-        tzinfo_factory = ((cursorObject *)curs)->tzinfo_factory;
-        if (n >= 5 && tzinfo_factory != Py_None) {
-            /* we have a time zone, calculate minutes and create
-               appropriate tzinfo object calling the factory */
-            Dprintf("typecast_PYDATETIME_cast: UTC offset = %ds", tz);
-
-            /* The datetime module requires that time zone offsets be
-               a whole number of minutes, so truncate the seconds to the
-               closest minute. */
-            // printf("%d %d %d\n", tz, tzmin, round(tz / 60.0));
-            tzinfo = PyObject_CallFunction(tzinfo_factory, "i",
-                (int)round(tz / 60.0));
-        } else {
-            Py_INCREF(Py_None);
-            tzinfo = Py_None;
-        }
-        if (tzinfo != NULL) {
-            obj = PyObject_CallFunction(
-                (PyObject*)PyDateTimeAPI->DateTimeType, "iiiiiiiO",
-                y, m, d, hh, mm, ss, us, tzinfo);
-            Dprintf("typecast_PYDATETIME_cast: tzinfo: %p, refcnt = "
-                FORMAT_CODE_PY_SSIZE_T,
-                tzinfo, Py_REFCNT(tzinfo)
-              );
-            Py_DECREF(tzinfo);
-        }
+static PyObject *
+typecast_PYDATETIMETZ_cast(const char *str, Py_ssize_t len, PyObject *curs)
+{
+    if (str == NULL) { Py_RETURN_NONE; }
+
+    if (!strcmp(str, "infinity") || !strcmp(str, "-infinity")) {
+        return _parse_inftz(str, curs);
     }
-    return obj;
+
+    return _parse_noninftz(str, len, curs);
 }
 
 /** TIME - parse time into a time object **/
@@ -345,4 +406,5 @@ typecast_PYINTERVAL_cast(const char *str, Py_ssize_t len, PyObject *curs)
 #define typecast_TIME_cast typecast_PYTIME_cast
 #define typecast_INTERVAL_cast typecast_PYINTERVAL_cast
 #define typecast_DATETIME_cast typecast_PYDATETIME_cast
+#define typecast_DATETIMETZ_cast typecast_PYDATETIMETZ_cast
 #endif
diff --git a/psycopg/typecast_mxdatetime.c b/psycopg/typecast_mxdatetime.c
index 4b03d15..12c734a 100644
--- a/psycopg/typecast_mxdatetime.c
+++ b/psycopg/typecast_mxdatetime.c
@@ -248,5 +248,6 @@ typecast_MXINTERVAL_cast(const char *str, Py_ssize_t len, PyObject *curs)
 #define typecast_TIME_cast typecast_MXTIME_cast
 #define typecast_INTERVAL_cast typecast_MXINTERVAL_cast
 #define typecast_DATETIME_cast typecast_MXDATE_cast
+#define typecast_DATETIMETZ_cast typecast_MXDATE_cast
 #endif
 
diff --git a/tests/test_dates.py b/tests/test_dates.py
index 98b1f55..83eea32 100755
--- a/tests/test_dates.py
+++ b/tests/test_dates.py
@@ -392,6 +392,25 @@ class DatetimeTests(ConnectingTestCase, CommonDatetimeTestsMixin):
         self.assertRaises(OverflowError, f, '00:00:100000000000000000:00')
         self.assertRaises(OverflowError, f, '00:00:00.100000000000000000')
 
+    def test_adapt_infinity_tz(self):
+        from datetime import datetime
+
+        t = self.execute("select 'infinity'::timestamp")
+        self.assert_(t.tzinfo is None)
+        self.assert_(t > datetime(4000, 1, 1))
+
+        t = self.execute("select '-infinity'::timestamp")
+        self.assert_(t.tzinfo is None)
+        self.assert_(t < datetime(1000, 1, 1))
+
+        t = self.execute("select 'infinity'::timestamptz")
+        self.assert_(t.tzinfo is not None)
+        self.assert_(t > datetime(4000, 1, 1, tzinfo=FixedOffsetTimezone()))
+
+        t = self.execute("select '-infinity'::timestamptz")
+        self.assert_(t.tzinfo is not None)
+        self.assert_(t < datetime(1000, 1, 1, tzinfo=FixedOffsetTimezone()))
+
 
 # Only run the datetime tests if psycopg was compiled with support.
 if not hasattr(psycopg2.extensions, 'PYDATETIME'):