Fixed time zone handling

timtroendle · timtroendle · commit 2081fc9e8eeb · 2014-11-22T20:17:44.000+01:00
diff --git a/influxdb/misc.py b/influxdb/misc.py
@@ -2,7 +2,6 @@
 """
 Miscellaneous
 """
-from time import mktime
 
 from .client import InfluxDBClient
 
@@ -14,6 +13,9 @@ class DataFrameClient(InfluxDBClient):
     The client reads and writes from pandas DataFrames.
     """
 
+    import pandas as pd
+    EPOCH = pd.Timestamp('1970-01-01 00:00:00.000+00:00')
+
     def write_points(self, data, *args, **kwargs):
         """
         write_points()
@@ -84,8 +86,13 @@ def _convert_dataframe_to_json(self, dataframe, name):
             raise TypeError('Must be DataFrame with DatetimeIndex or \
                             PeriodIndex.')
         dataframe.index = dataframe.index.to_datetime()
-        dataframe['time'] = [mktime(dt.timetuple()) for dt in dataframe.index]
+        if dataframe.index.tzinfo == None:
+            dataframe.index = dataframe.index.tz_localize('UTC')
+        dataframe['time'] = [self._datetime_to_epoch(dt) for dt in dataframe.index]
         data = {'name': name,
                 'columns': [str(column) for column in dataframe.columns],
                 'points': list([list(x) for x in dataframe.values])}
         return data
+
+    def _datetime_to_epoch(self, datetime):
+        return (datetime - DataFrameClient.EPOCH).total_seconds()
diff --git a/tests/influxdb/misc_test.py b/tests/influxdb/misc_test.py
@@ -6,8 +6,7 @@
 import json
 import requests_mock
 from nose.tools import raises
-from datetime import datetime, timedelta
-import time
+from datetime import timedelta
 from tests import skipIfPYpy, using_pypy
 
 if not using_pypy:
@@ -22,17 +21,16 @@
 class TestDataFrameClient(unittest.TestCase):
 
     def test_write_points_from_dataframe(self):
-        now = datetime(2014, 11, 15, 15, 42, 44, 543)
+        now = pd.Timestamp('1970-01-01 00:00+00:00')
         dataframe = pd.DataFrame(data=[["1", 1, 1.0], ["2", 2, 2.0]],
                                  index=[now, now + timedelta(hours=1)],
                                  columns=["column_one", "column_two",
                                           "column_three"])
         points = [
             {
                 "points": [
-                    ["1", 1, 1.0, time.mktime(now.timetuple())],
-                    ["2", 2, 2.0, time.mktime((now + timedelta(hours=1))
-                                              .timetuple())]
+                    ["1", 1, 1.0, 0],
+                    ["2", 2, 2.0, 3600]
                 ],
                 "name": "foo",
                 "columns": ["column_one", "column_two", "column_three", "time"]
@@ -49,16 +47,15 @@ def test_write_points_from_dataframe(self):
             self.assertListEqual(json.loads(m.last_request.body), points)
 
     def test_write_points_from_dataframe_with_numeric_column_names(self):
-        now = datetime(2014, 11, 15, 15, 42, 44, 543)
+        now = pd.Timestamp('1970-01-01 00:00+00:00')
         # df with numeric column names
         dataframe = pd.DataFrame(data=[["1", 1, 1.0], ["2", 2, 2.0]],
                                  index=[now, now + timedelta(hours=1)])
         points = [
             {
                 "points": [
-                    ["1", 1, 1.0, time.mktime(now.timetuple())],
-                    ["2", 2, 2.0, time.mktime((now + timedelta(hours=1))
-                                              .timetuple())]
+                    ["1", 1, 1.0, 0],
+                    ["2", 2, 2.0, 3600]
                 ],
                 "name": "foo",
                 "columns": ['0', '1', '2', "time"]
@@ -75,18 +72,16 @@ def test_write_points_from_dataframe_with_numeric_column_names(self):
             self.assertListEqual(json.loads(m.last_request.body), points)
 
     def test_write_points_from_dataframe_with_period_index(self):
-        now = datetime(2014, 11, 16)
         dataframe = pd.DataFrame(data=[["1", 1, 1.0], ["2", 2, 2.0]],
-                                 index=[pd.Period('2014-11-16'),
-                                        pd.Period('2014-11-17')],
+                                 index=[pd.Period('1970-01-01'),
+                                        pd.Period('1970-01-02')],
                                  columns=["column_one", "column_two",
                                           "column_three"])
         points = [
             {
                 "points": [
-                    ["1", 1, 1.0, time.mktime(now.timetuple())],
-                    ["2", 2, 2.0, time.mktime((now + timedelta(hours=24))
-                                              .timetuple())]
+                    ["1", 1, 1.0, 0],
+                    ["2", 2, 2.0, 86400]
                 ],
                 "name": "foo",
                 "columns": ["column_one", "column_two", "column_three", "time"]
@@ -117,7 +112,7 @@ def test_write_points_from_dataframe_fails_without_time_index(self):
 
     @raises(TypeError)
     def test_write_points_from_dataframe_fails_with_series(self):
-        now = datetime(2014, 11, 16)
+        now = pd.Timestamp('1970-01-01 00:00+00:00')
         dataframe = pd.Series(data=[1.0, 2.0],
                               index=[now, now + timedelta(hours=1)])
 
