Cast datetime into timestamptz when the tzinfo field is set.

2 files changed, 17 insertions, 2 deletions

diff --git a/psycopg/adapter_datetime.c b/psycopg/adapter_datetime.c
index deaf59f..24d158c 100644
--- a/psycopg/adapter_datetime.c
+++ b/psycopg/adapter_datetime.c
@@ -57,6 +57,7 @@ static PyObject *
 pydatetime_str(pydatetimeObject *self)
 {
     if (self->type <= PSYCO_DATETIME_TIMESTAMP) {
+        PyObject *tz;
 
         /* Select the right PG type to cast into. */
         char *fmt = NULL;
@@ -68,7 +69,10 @@ pydatetime_str(pydatetimeObject *self)
             fmt = "'%s'::date";
             break;
         case PSYCO_DATETIME_TIMESTAMP:
-            fmt = "'%s'::timestamp";
+            tz = PyObject_GetAttrString(self->wrapped, "tzinfo");
+            if (!tz) { return NULL; }
+            fmt = (tz == Py_None) ? "'%s'::timestamp" : "'%s'::timestamptz";
+            Py_DECREF(tz);
             break;
         }
 
diff --git a/tests/test_dates.py b/tests/test_dates.py
index 996d591..1ffb15a 100644
--- a/tests/test_dates.py
+++ b/tests/test_dates.py
@@ -250,6 +250,7 @@ class DatetimeTests(unittest.TestCase, CommonDatetimeTestsMixin):
     def _test_type_roundtrip(self, o1):
         o2 = self.execute("select %s;", (o1,))
         self.assertEqual(type(o1), type(o2))
+        return o2
 
     def _test_type_roundtrip_array(self, o1):
         o1 = [o1]
@@ -262,7 +263,17 @@ class DatetimeTests(unittest.TestCase, CommonDatetimeTestsMixin):
 
     def test_type_roundtrip_datetime(self):
         from datetime import datetime
-        self._test_type_roundtrip(datetime(2010,05,03,10,20,30))
+        dt = self._test_type_roundtrip(datetime(2010,05,03,10,20,30))
+        self.assertEqual(None, dt.tzinfo)
+
+    def test_type_roundtrip_datetimetz(self):
+        from datetime import datetime
+        import psycopg2.tz
+        tz = psycopg2.tz.FixedOffsetTimezone(8*60)
+        dt1 = datetime(2010,05,03,10,20,30, tzinfo=tz)
+        dt2 = self._test_type_roundtrip(dt1)
+        self.assertNotEqual(None, dt2.tzinfo)
+        self.assertEqual(dt1, dt2)
 
     def test_type_roundtrip_time(self):
         from datetime import time