Fix timedeltas to work with to_json

Author: cpcloud

pandas/io/tests/test_json/test_pandas.py

@@ -4,8 +4,8 @@
 import os
 
 import numpy as np
-import nose
 from pandas import Series, DataFrame, DatetimeIndex, Timestamp
+from datetime import timedelta
 import pandas as pd
 read_json = pd.read_json
 
@@ -601,7 +601,6 @@ def test_url(self):
             self.assertEqual(result[c].dtype, 'datetime64[ns]')
 
     def test_timedelta(self):
-        from datetime import timedelta
         converter = lambda x: pd.to_timedelta(x,unit='ms')
 
         s = Series([timedelta(23), timedelta(seconds=5)])
@@ -613,17 +612,32 @@ def test_timedelta(self):
         assert_frame_equal(
             frame, pd.read_json(frame.to_json()).apply(converter))
 
-    def test_default_handler(self):
-        from datetime import timedelta
-
-        frame = DataFrame([timedelta(23), timedelta(seconds=5), 42])
-        self.assertRaises(OverflowError, frame.to_json)
+        frame = DataFrame({'a': [timedelta(23), timedelta(seconds=5)],
+                           'b': [1, 2],
+                           'c': pd.date_range(start='20130101', periods=2)})
+        result = pd.read_json(frame.to_json(date_unit='ns'))
+        result['a'] = pd.to_timedelta(result.a, unit='ns')
+        result['c'] = pd.to_datetime(result.c)
+        assert_frame_equal(frame, result)
+
+    def test_mixed_timedelta_datetime(self):
+        frame = DataFrame({'a': [timedelta(23), pd.Timestamp('20130101')]},
+                          dtype=object)
+        expected = pd.read_json(frame.to_json(date_unit='ns'),
+                                dtype={'a': 'int64'})
+        assert_frame_equal(DataFrame({'a': [pd.Timedelta(frame.a[0]).value,
+                                            pd.Timestamp(frame.a[1]).value]}),
+                           expected)
 
-        expected = DataFrame([str(timedelta(23)), str(timedelta(seconds=5)), 42])
-        assert_frame_equal(
-            expected, pd.read_json(frame.to_json(default_handler=str)))
+    def test_default_handler(self):
+        value = object()
+        frame = DataFrame({'a': ['a', value]})
+        expected = frame.applymap(str)
+        result = pd.read_json(frame.to_json(default_handler=str))
+        assert_frame_equal(expected, result)
 
+    def test_default_handler_raises(self):
         def my_handler_raises(obj):
             raise TypeError("raisin")
-        self.assertRaises(TypeError, frame.to_json,
+        self.assertRaises(TypeError, DataFrame({'a': [1, 2, object()]}).to_json,
                           default_handler=my_handler_raises)

pandas/src/datetime_helper.h

@@ -1,6 +1,23 @@
 #include "datetime.h"
 
+#if PY_MAJOR_VERSION >= 3
+#define PyInt_AS_LONG PyLong_AsLong
+#endif
+
 void mangle_nat(PyObject *val) {
   PyDateTime_GET_MONTH(val) = -1;
   PyDateTime_GET_DAY(val) = -1;
 }
+
+long get_long_attr(PyObject *o, const char *attr) {
+  return PyInt_AS_LONG(PyObject_GetAttrString(o, attr));
+}
+
+double total_seconds(PyObject *td) {
+  // Python 2.6 compat
+  long microseconds = get_long_attr(td, "microseconds");
+  long seconds = get_long_attr(td, "seconds");
+  long days = get_long_attr(td, "days");
+  long days_in_seconds = days * 24 * 3600;
+  return (microseconds + (seconds + days_in_seconds) * 1000000.0) / 1000000.0;
+}

pandas/src/ujson/python/objToJSON.c

@@ -41,11 +41,11 @@ Numeric decoder derived from from TCL library
 #include <numpy/arrayscalars.h>
 #include <np_datetime.h>
 #include <np_datetime_strings.h>
+#include <datetime_helper.h>
 #include <numpy_helper.h>
 #include <numpy/npy_math.h>
 #include <math.h>
 #include <stdio.h>
-#include <datetime.h>
 #include <ultrajson.h>
 
 static PyObject* type_decimal;
@@ -154,6 +154,7 @@ enum PANDAS_FORMAT
 // import_array() compat
 #if (PY_VERSION_HEX >= 0x03000000)
 void *initObjToJSON(void)
+
 #else
 void initObjToJSON(void)
 #endif
@@ -1447,7 +1448,7 @@ void Object_beginTypeContext (JSOBJ _obj, JSONTypeContext *tc)
     pc->PyTypeToJSON = NpyDateTimeToJSON;
     if (enc->datetimeIso)
     {
-      tc->type = JT_UTF8;
+  tc->type = JT_UTF8;
     }
     else
     {
@@ -1456,6 +1457,37 @@ void Object_beginTypeContext (JSOBJ _obj, JSONTypeContext *tc)
     return;
   }
   else
+  if (PyDelta_Check(obj))
+  {
+    long value;
+
+    if (PyObject_HasAttrString(obj, "value"))
+        value = get_long_attr(obj, "value");
+    else
+        value = total_seconds(obj) * 1000000000; // nanoseconds per second
+
+    exc = PyErr_Occurred();
+
+    if (exc && PyErr_ExceptionMatches(PyExc_OverflowError))
+    {
+      PRINTMARK();
+      goto INVALID;
+    }
+
+    if (value == get_nat()) {
+      PRINTMARK();
+      tc->type = JT_NULL;
+      return;
+    }
+
+    GET_TC(tc)->longValue = value;
+
+    PRINTMARK();
+    pc->PyTypeToJSON = PyLongToINT64;
+    tc->type = JT_LONG;
+    return;
+  }
+  else
   if (PyArray_IsScalar(obj, Integer))
   {
     PRINTMARK();

pandas/tslib.pyx

@@ -20,6 +20,9 @@ cdef extern from "Python.h":
     cdef PyTypeObject *Py_TYPE(object)
     int PySlice_Check(object)
 
+cdef extern from "datetime_helper.h":
+    double total_seconds(object)
+
 # this is our datetime.pxd
 from datetime cimport *
 from util cimport is_integer_object, is_float_object, is_datetime64_object, is_timedelta64_object
@@ -2753,10 +2756,6 @@ cdef object _get_deltas(object tz):
 
     return utc_offset_cache[cache_key]
 
-cdef double total_seconds(object td): # Python 2.6 compat
-    return ((td.microseconds + (td.seconds + td.days * 24 * 3600) * 10**6) //
-            10**6)
-
 def tot_seconds(td):
     return total_seconds(td)