ENH: datetime-string: Add support for datetime.tzinfo in datetime_as_string

Also add some tests using the pytz library.

6 files changed, 340 insertions, 185 deletions

diff --git a/numpy/core/src/multiarray/_datetime.h b/numpy/core/src/multiarray/_datetime.h
index e11e0eeeb..be9e6978f 100644
--- a/numpy/core/src/multiarray/_datetime.h
+++ b/numpy/core/src/multiarray/_datetime.h
@@ -211,6 +211,13 @@ NPY_NO_EXPORT PyObject *
 convert_datetime_metadata_tuple_to_metacobj(PyObject *tuple);
 
 /*
+ * Gets a tzoffset in minutes by calling the fromutc() function on
+ * the Python datetime.tzinfo object.
+ */
+NPY_NO_EXPORT int
+get_tzoffset_from_pytzinfo(PyObject *timezone, npy_datetimestruct *dts);
+
+/*
  * Converts an input object into datetime metadata. The input
  * may be either a string or a tuple.
  *
@@ -240,12 +247,16 @@ append_metastr_to_string(PyArray_DatetimeMetaData *meta,
  * 'out_bestunit' gives a suggested unit based on whether the object
  *      was a datetime.date or datetime.datetime object.
  *
+ * If 'apply_tzinfo' is 1, this function uses the tzinfo to convert
+ * to UTC time, otherwise it returns the struct with the local time.
+ *
  * Returns -1 on error, 0 on success, and 1 (with no error set)
  * if obj doesn't have the neeeded date or datetime attributes.
  */
 NPY_NO_EXPORT int
 convert_pydatetime_to_datetimestruct(PyObject *obj, npy_datetimestruct *out,
-                                     NPY_DATETIMEUNIT *out_bestunit);
+                                     NPY_DATETIMEUNIT *out_bestunit,
+                                     int apply_tzinfo);
 
 /*
  * Converts a PyObject * into a datetime, in any of the forms supported.
diff --git a/numpy/core/src/multiarray/datetime.c b/numpy/core/src/multiarray/datetime.c
index 301677895..77c7c9605 100644
--- a/numpy/core/src/multiarray/datetime.c
+++ b/numpy/core/src/multiarray/datetime.c
@@ -131,6 +131,19 @@ get_datetimestruct_days(const npy_datetimestruct *dts)
 }
 
 /*
+ * Calculates the minutes offset from the 1970 epoch.
+ */
+NPY_NO_EXPORT npy_int64
+get_datetimestruct_minutes(const npy_datetimestruct *dts)
+{
+    npy_int64 days = get_datetimestruct_days(dts) * 24 * 60;
+    days += dts->hour * 60;
+    days += dts->min;
+
+    return days;
+}
+
+/*
  * Modifies '*days_' to be the day offset within the year,
  * and returns the year.
  */
@@ -2451,12 +2464,16 @@ add_minutes_to_datetimestruct(npy_datetimestruct *dts, int minutes)
  * 'out_bestunit' gives a suggested unit based on whether the object
  *      was a datetime.date or datetime.datetime object.
  *
+ * If 'apply_tzinfo' is 1, this function uses the tzinfo to convert
+ * to UTC time, otherwise it returns the struct with the local time.
+ *
  * Returns -1 on error, 0 on success, and 1 (with no error set)
  * if obj doesn't have the neeeded date or datetime attributes.
  */
 NPY_NO_EXPORT int
 convert_pydatetime_to_datetimestruct(PyObject *obj, npy_datetimestruct *out,
-                                     NPY_DATETIMEUNIT *out_bestunit)
+                                     NPY_DATETIMEUNIT *out_bestunit,
+                                     int apply_tzinfo)
 {
     PyObject *tmp;
     int isleap;
@@ -2514,7 +2531,8 @@ convert_pydatetime_to_datetimestruct(PyObject *obj, npy_datetimestruct *out,
         goto invalid_date;
     }
     isleap = is_leapyear(out->year);
-    if (out->day < 1 || out->day > _days_per_month_table[isleap][out->month-1]) {
+    if (out->day < 1 ||
+                out->day > _days_per_month_table[isleap][out->month-1]) {
         goto invalid_date;
     }
 
@@ -2586,7 +2604,7 @@ convert_pydatetime_to_datetimestruct(PyObject *obj, npy_datetimestruct *out,
     }
 
     /* Apply the time zone offset if it exists */
-    if (PyObject_HasAttrString(obj, "tzinfo")) {
+    if (apply_tzinfo && PyObject_HasAttrString(obj, "tzinfo")) {
         tmp = PyObject_GetAttrString(obj, "tzinfo");
         if (tmp == NULL) {
             return -1;
@@ -2607,10 +2625,10 @@ convert_pydatetime_to_datetimestruct(PyObject *obj, npy_datetimestruct *out,
             Py_DECREF(tmp);
 
             /*
-             * The timedelta should have an attribute "seconds"
+             * The timedelta should have a function "total_seconds"
              * which contains the value we want.
              */
-            tmp = PyObject_GetAttrString(obj, "seconds");
+            tmp = PyObject_CallMethod(offset, "total_seconds", "");
             if (tmp == NULL) {
                 return -1;
             }
@@ -2650,6 +2668,43 @@ invalid_time:
 }
 
 /*
+ * Gets a tzoffset in minutes by calling the fromutc() function on
+ * the Python datetime.tzinfo object.
+ */
+NPY_NO_EXPORT int
+get_tzoffset_from_pytzinfo(PyObject *timezone, npy_datetimestruct *dts)
+{
+    PyObject *dt, *loc_dt;
+    npy_datetimestruct loc_dts;
+
+    /* Create a Python datetime to give to the timezone object */
+    dt = PyDateTime_FromDateAndTime((int)dts->year, dts->month, dts->day,
+                            dts->hour, dts->min, 0, 0);
+    if (dt == NULL) {
+        return -1;
+    }
+
+    /* Convert the datetime from UTC to local time */
+    loc_dt = PyObject_CallMethod(timezone, "fromutc", "O", dt);
+    Py_DECREF(dt);
+    if (loc_dt == NULL) {
+        return -1;
+    }
+
+    /* Convert the local datetime into a datetimestruct */
+    if (convert_pydatetime_to_datetimestruct(loc_dt, &loc_dts, NULL, 0) < 0) {
+        Py_DECREF(loc_dt);
+        return -1;
+    }
+
+    Py_DECREF(loc_dt);
+
+    /* Calculate the tzoffset as the difference between the datetimes */
+    return get_datetimestruct_minutes(&loc_dts) -
+           get_datetimestruct_minutes(dts);
+}
+
+/*
  * Converts a PyObject * into a datetime, in any of the forms supported.
  *
  * If the units metadata isn't known ahead of time, set meta->base
@@ -2793,7 +2848,7 @@ convert_pyobject_to_datetime(PyArray_DatetimeMetaData *meta, PyObject *obj,
         npy_datetimestruct dts;
         NPY_DATETIMEUNIT bestunit = -1;
 
-        code = convert_pydatetime_to_datetimestruct(obj, &dts, &bestunit);
+        code = convert_pydatetime_to_datetimestruct(obj, &dts, &bestunit, 1);
         if (code == -1) {
             return -1;
         }
diff --git a/numpy/core/src/multiarray/datetime_strings.c b/numpy/core/src/multiarray/datetime_strings.c
index 679906aee..8310a80fe 100644
--- a/numpy/core/src/multiarray/datetime_strings.c
+++ b/numpy/core/src/multiarray/datetime_strings.c
@@ -834,11 +834,6 @@ make_iso_8601_datetime(npy_datetimestruct *dts, char *outstr, int outlen,
         base = NPY_FR_D;
     }
 
-    /* Printed dates have no time zone */
-    if (base < NPY_FR_h) {
-        local = 0;
-    }
-
     /* Use the C API to convert from UTC to local time */
     if (local && tzoffset == -1) {
         time_t rawtime = 0, localrawtime;
@@ -1217,3 +1212,226 @@ string_too_short:
 }
 
 
+/*
+ * This is the Python-exposed datetime_as_string function.
+ */
+NPY_NO_EXPORT PyObject *
+array_datetime_as_string(PyObject *NPY_UNUSED(self), PyObject *args,
+                                PyObject *kwds)
+{
+    PyObject *arr_in = NULL, *unit_in = NULL, *timezone = NULL;
+
+    int local = 0;
+    NPY_DATETIMEUNIT unit;
+    PyArray_DatetimeMetaData *meta;
+    int strsize;
+
+    PyArrayObject *ret = NULL;
+
+    NpyIter *iter = NULL;
+    PyArrayObject *op[2] = {NULL, NULL};
+    PyArray_Descr *op_dtypes[2] = {NULL, NULL};
+    npy_uint32 flags, op_flags[2];
+
+    static char *kwlist[] = {"arr", "unit", "timezone", NULL};
+
+    if(!PyArg_ParseTupleAndKeywords(args, kwds,
+                                "O|OO:datetime_as_string", kwlist,
+                                &arr_in,
+                                &unit_in,
+                                &timezone)) {
+        return NULL;
+    }
+
+    /* Claim a reference to timezone for later */
+    Py_XINCREF(timezone);
+
+    op[0] = (PyArrayObject *)PyArray_FromAny(arr_in,
+                                    NULL, 0, 0, 0, NULL);
+    if (PyArray_DESCR(op[0])->type_num != NPY_DATETIME) {
+        PyErr_SetString(PyExc_TypeError,
+                    "input must have type NumPy datetime");
+        goto fail;
+    }
+
+    /* Get the datetime metadata */
+    meta = get_datetime_metadata_from_dtype(PyArray_DESCR(op[0]));
+    if (meta == NULL) {
+        goto fail;
+    }
+
+    /* Use the metadata's unit for printing by default */
+    unit = meta->base;
+
+    /* Parse the input unit if provided */
+    if (unit_in != NULL && unit_in != Py_None) {
+        PyObject *strobj;
+        char *str = NULL;
+        Py_ssize_t len = 0;
+
+        if (PyUnicode_Check(unit_in)) {
+            strobj = PyUnicode_AsASCIIString(unit_in);
+            if (strobj == NULL) {
+                goto fail;
+            }
+        }
+        else {
+            strobj = unit_in;
+            Py_INCREF(strobj);
+        }
+
+        if (PyBytes_AsStringAndSize(strobj, &str, &len) < 0) {
+            Py_DECREF(strobj);
+            goto fail;
+        }
+
+        /* unit == -1 means to autodetect the unit from the datetime data */
+        if (strcmp(str, "auto") == 0) {
+            unit = -1;
+        }
+        else {
+            unit = parse_datetime_unit_from_string(str, len, NULL);
+            if (unit == -1) {
+                Py_DECREF(strobj);
+                goto fail;
+            }
+        }
+        Py_DECREF(strobj);
+    }
+
+    /* Get the input time zone */
+    if (timezone != NULL) {
+        /* Convert to ASCII if it's unicode */
+        if (PyUnicode_Check(timezone)) {
+            /* accept unicode input */
+            PyObject *obj_str;
+            obj_str = PyUnicode_AsASCIIString(timezone);
+            if (obj_str == NULL) {
+                goto fail;
+            }
+            Py_DECREF(timezone);
+            timezone = obj_str;
+        }
+
+        /* Check for the supported string inputs */
+        if (PyBytes_Check(timezone)) {
+            char *str;
+            Py_ssize_t len;
+
+            if (PyBytes_AsStringAndSize(timezone, &str, &len) < 0) {
+                goto fail;
+            }
+
+            if (strcmp(str, "local") == 0) {
+                local = 1;
+                Py_DECREF(timezone);
+                timezone = NULL;
+            }
+            else if (strcmp(str, "UTC") == 0) {
+                local = 0;
+                Py_DECREF(timezone);
+                timezone = NULL;
+            }
+            else {
+                PyErr_Format(PyExc_ValueError, "Unsupported timezone "
+                            "input string \"%s\"", str);
+                goto fail;
+            }
+        }
+        /* Otherwise assume it's a Python TZInfo, or acts like one */
+        else {
+            local = 1;
+        }
+    }
+
+    /* Create the output string data type with a big enough length */
+    op_dtypes[1] = PyArray_DescrNewFromType(NPY_STRING);
+    if (op_dtypes[1] == NULL) {
+        goto fail;
+    }
+    strsize = get_datetime_iso_8601_strlen(local, unit);
+    op_dtypes[1]->elsize = strsize;
+
+    flags = NPY_ITER_ZEROSIZE_OK|
+            NPY_ITER_BUFFERED;
+    op_flags[0] = NPY_ITER_READONLY|
+                  NPY_ITER_ALIGNED;
+    op_flags[1] = NPY_ITER_WRITEONLY|
+                  NPY_ITER_ALLOCATE;
+
+    iter = NpyIter_MultiNew(2, op, flags, NPY_KEEPORDER, NPY_NO_CASTING,
+                            op_flags, op_dtypes);
+    if (iter == NULL) {
+        goto fail;
+    }
+
+    if (NpyIter_GetIterSize(iter) != 0) {
+        NpyIter_IterNextFunc *iternext;
+        char **dataptr;
+        npy_datetime dt;
+        npy_datetimestruct dts;
+
+        iternext = NpyIter_GetIterNext(iter, NULL);
+        if (iternext == NULL) {
+            goto fail;
+        }
+        dataptr = NpyIter_GetDataPtrArray(iter);
+
+        do {
+            int tzoffset = -1;
+
+            /* Get the datetime */
+            dt = *(npy_datetime *)dataptr[0];
+
+            /* Convert it to a struct */
+            if (convert_datetime_to_datetimestruct(meta, dt, &dts) < 0) {
+                goto fail;
+            }
+
+            /* Get the tzoffset from the timezone if provided */
+            if (local && timezone != NULL) {
+                tzoffset = get_tzoffset_from_pytzinfo(timezone, &dts);
+                if (tzoffset == -1) {
+                    goto fail;
+                }
+            }
+
+            /* Zero the destination string completely */
+            memset(dataptr[1], 0, strsize);
+            /* Convert that into a string */
+            if (make_iso_8601_datetime(&dts, (char *)dataptr[1], strsize,
+                                local, unit, tzoffset) < 0) {
+                goto fail;
+            }
+        } while(iternext(iter));
+    }
+
+    ret = NpyIter_GetOperandArray(iter)[1];
+    Py_INCREF(ret);
+
+    Py_XDECREF(timezone);
+    Py_XDECREF(op[0]);
+    Py_XDECREF(op[1]);
+    Py_XDECREF(op_dtypes[0]);
+    Py_XDECREF(op_dtypes[1]);
+    if (iter != NULL) {
+        NpyIter_Deallocate(iter);
+    }
+
+    return PyArray_Return(ret);
+
+fail:
+    Py_XDECREF(timezone);
+    Py_XDECREF(op[0]);
+    Py_XDECREF(op[1]);
+    Py_XDECREF(op_dtypes[0]);
+    Py_XDECREF(op_dtypes[1]);
+    if (iter != NULL) {
+        NpyIter_Deallocate(iter);
+    }
+
+    return NULL;
+}
+
+
+
diff --git a/numpy/core/src/multiarray/datetime_strings.h b/numpy/core/src/multiarray/datetime_strings.h
index 2b48f49b9..e855a8631 100644
--- a/numpy/core/src/multiarray/datetime_strings.h
+++ b/numpy/core/src/multiarray/datetime_strings.h
@@ -74,4 +74,11 @@ NPY_NO_EXPORT int
 make_iso_8601_datetime(npy_datetimestruct *dts, char *outstr, int outlen,
                     int local, NPY_DATETIMEUNIT base, int tzoffset);
 
+/*
+ * This is the Python-exposed datetime_as_string function.
+ */
+NPY_NO_EXPORT PyObject *
+array_datetime_as_string(PyObject *NPY_UNUSED(self), PyObject *args,
+                                PyObject *kwds);
+
 #endif
diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c
index 72b1a5a65..3b537a556 100644
--- a/numpy/core/src/multiarray/multiarraymodule.c
+++ b/numpy/core/src/multiarray/multiarraymodule.c
@@ -2809,170 +2809,6 @@ array_datetime_data(PyObject *NPY_UNUSED(dummy), PyObject *args)
     return convert_datetime_metadata_to_tuple(meta);
 }
 
-static PyObject *
-array_datetime_as_string(PyObject *NPY_UNUSED(self), PyObject *args,
-                                PyObject *kwds)
-{
-    PyObject *arr_in = NULL, *unit_in = NULL;
-    int local = 0, tzoffset = -1;
-    NPY_DATETIMEUNIT unit;
-    PyArray_DatetimeMetaData *meta;
-    int strsize;
-
-    PyArrayObject *ret = NULL;
-
-    NpyIter *iter = NULL;
-    PyArrayObject *op[2] = {NULL, NULL};
-    PyArray_Descr *op_dtypes[2] = {NULL, NULL};
-    npy_uint32 flags, op_flags[2];
-
-    static char *kwlist[] = {"arr", "local", "unit", "tzoffset", NULL};
-
-    if(!PyArg_ParseTupleAndKeywords(args, kwds,
-                                "O|iOi:datetime_as_string", kwlist,
-                                &arr_in,
-                                &local,
-                                &unit_in,
-                                &tzoffset)) {
-        goto fail;
-    }
-
-    op[0] = (PyArrayObject *)PyArray_FromAny(arr_in,
-                                    NULL, 0, 0, 0, NULL);
-    if (PyArray_DESCR(op[0])->type_num != NPY_DATETIME) {
-        PyErr_SetString(PyExc_TypeError,
-                    "input must have type NumPy datetime");
-        goto fail;
-    }
-
-    /* Get the datetime metadata */
-    meta = get_datetime_metadata_from_dtype(PyArray_DESCR(op[0]));
-    if (meta == NULL) {
-        goto fail;
-    }
-
-    /* Use the metadata's unit for printing by default */
-    unit = meta->base;
-
-    /* Parse the input unit if provided */
-    if (unit_in != NULL && unit_in != Py_None) {
-        PyObject *strobj;
-        char *str = NULL;
-        Py_ssize_t len = 0;
-
-        if (PyUnicode_Check(unit_in)) {
-            strobj = PyUnicode_AsASCIIString(unit_in);
-            if (strobj == NULL) {
-                goto fail;
-            }
-        }
-        else {
-            strobj = unit_in;
-            Py_INCREF(strobj);
-        }
-
-        if (PyBytes_AsStringAndSize(strobj, &str, &len) < 0) {
-            Py_DECREF(strobj);
-            goto fail;
-        }
-
-        /* unit == -1 means to autodetect the unit from the datetime data */
-        if (strcmp(str, "auto") == 0) {
-            unit = -1;
-        }
-        else {
-            unit = parse_datetime_unit_from_string(str, len, NULL);
-            if (unit == -1) {
-                Py_DECREF(strobj);
-                goto fail;
-            }
-        }
-        Py_DECREF(strobj);
-    }
-
-    if (!local && tzoffset != -1) {
-        PyErr_SetString(PyExc_ValueError,
-                "Can only use 'tzoffset' parameter when 'local' is "
-                "set to True");
-        goto fail;
-    }
-
-    /* Create the output string data type with a big enough length */
-    op_dtypes[1] = PyArray_DescrNewFromType(NPY_STRING);
-    if (op_dtypes[1] == NULL) {
-        goto fail;
-    }
-    strsize = get_datetime_iso_8601_strlen(local, unit);
-    op_dtypes[1]->elsize = strsize;
-
-    flags = NPY_ITER_ZEROSIZE_OK|
-            NPY_ITER_BUFFERED;
-    op_flags[0] = NPY_ITER_READONLY|
-                  NPY_ITER_ALIGNED;
-    op_flags[1] = NPY_ITER_WRITEONLY|
-                  NPY_ITER_ALLOCATE;
-
-    iter = NpyIter_MultiNew(2, op, flags, NPY_KEEPORDER, NPY_NO_CASTING,
-                            op_flags, op_dtypes);
-    if (iter == NULL) {
-        goto fail;
-    }
-
-    if (NpyIter_GetIterSize(iter) != 0) {
-        NpyIter_IterNextFunc *iternext;
-        char **dataptr;
-        npy_datetime dt;
-        npy_datetimestruct dts;
-
-        iternext = NpyIter_GetIterNext(iter, NULL);
-        if (iternext == NULL) {
-            goto fail;
-        }
-        dataptr = NpyIter_GetDataPtrArray(iter);
-
-        do {
-            /* Get the datetime */
-            dt = *(datetime *)dataptr[0];
-            /* Convert it to a struct */
-            if (convert_datetime_to_datetimestruct(meta, dt, &dts) < 0) {
-                goto fail;
-            }
-            /* Zero the destination string completely */
-            memset(dataptr[1], 0, strsize);
-            /* Convert that into a string */
-            if (make_iso_8601_datetime(&dts, (char *)dataptr[1], strsize,
-                                local, unit, tzoffset) < 0) {
-                goto fail;
-            }
-        } while(iternext(iter));
-    }
-
-    ret = NpyIter_GetOperandArray(iter)[1];
-    Py_INCREF(ret);
-
-    Py_XDECREF(op[0]);
-    Py_XDECREF(op[1]);
-    Py_XDECREF(op_dtypes[0]);
-    Py_XDECREF(op_dtypes[1]);
-    if (iter != NULL) {
-        NpyIter_Deallocate(iter);
-    }
-
-    return PyArray_Return(ret);
-
-fail:
-    Py_XDECREF(op[0]);
-    Py_XDECREF(op[1]);
-    Py_XDECREF(op_dtypes[0]);
-    Py_XDECREF(op_dtypes[1]);
-    if (iter != NULL) {
-        NpyIter_Deallocate(iter);
-    }
-
-    return NULL;
-}
-
-
 #if !defined(NPY_PY3K)
 static PyObject *
 new_buffer(PyObject *NPY_UNUSED(dummy), PyObject *args)
diff --git a/numpy/core/tests/test_datetime.py b/numpy/core/tests/test_datetime.py
index 863cefea3..b198f6de6 100644
--- a/numpy/core/tests/test_datetime.py
+++ b/numpy/core/tests/test_datetime.py
@@ -1108,16 +1108,44 @@ class TestDateTime(TestCase):
                                             unit='auto'),
                             '2032-01-01')
 
-        # local=True
+    def test_datetime_as_string_timzone(self):
+        # timezone='local' vs 'UTC'
         a = np.datetime64('2010-03-15T06:30Z', 'm')
-        assert_(np.datetime_as_string(a, local=True) != '2010-03-15T6:30Z')
-        # local=True with tzoffset
-        assert_equal(np.datetime_as_string(a, local=True, tzoffset=-60),
-                     '2010-03-15T05:30-0100')
-        assert_equal(np.datetime_as_string(a, local=True, tzoffset=+30),
-                     '2010-03-15T07:00+0030')
-        assert_equal(np.datetime_as_string(a, local=True, tzoffset=-5*60),
-                     '2010-03-15T01:30-0500')
+        assert_equal(np.datetime_as_string(a, timezone='UTC'),
+                '2010-03-15T06:30Z')
+        assert_(np.datetime_as_string(a, timezone='local') !=
+                '2010-03-15T06:30Z')
+
+        # Use pytz to test out various time zones
+        try:
+            from pytz import timezone as tz
+
+            b = np.datetime64('2010-02-15T06:30Z', 'm')
+
+            assert_equal(np.datetime_as_string(a, timezone=tz('US/Central')),
+                         '2010-03-15T01:30-0500')
+            assert_equal(np.datetime_as_string(a, timezone=tz('US/Eastern')),
+                         '2010-03-15T02:30-0400')
+            assert_equal(np.datetime_as_string(a, timezone=tz('US/Pacific')),
+                         '2010-03-14T23:30-0700')
+
+            assert_equal(np.datetime_as_string(b, timezone=tz('US/Central')),
+                         '2010-02-15T00:30-0600')
+            assert_equal(np.datetime_as_string(b, timezone=tz('US/Eastern')),
+                         '2010-02-15T01:30-0500')
+            assert_equal(np.datetime_as_string(b, timezone=tz('US/Pacific')),
+                         '2010-02-14T22:30-0800')
+
+            # Check that we can print out the date in the specified time zone
+            assert_equal(np.datetime_as_string(a, unit='D',
+                                               timezone=tz('US/Pacific')),
+                         '2010-03-14')
+            assert_equal(np.datetime_as_string(b, unit='D',
+                                               timezone=tz('US/Central')),
+                         '2010-02-15')
+        except ImportError:
+            import warnings
+            warnings.warn("Need pytz library to test datetime timezones")
 
     def test_datetime_arange(self):
         # With two datetimes provided as strings