Merge pull request influxdata#93 from tzonghao/master

aviau · aviau · commit 0502fa22268e · 2014-12-11T11:10:00.000-05:00
Add time_precision option to datetime_to_epoch. (Thanks @tzonghao)
diff --git a/influxdb/misc.py b/influxdb/misc.py
@@ -32,6 +32,7 @@ def write_points(self, data, *args, **kwargs):
         """
 
         batch_size = kwargs.get('batch_size')
+        time_precision = kwargs.get('time_precision', 's')
         if batch_size:
             kwargs.pop('batch_size')  # don't hand over to InfluxDBClient
             for key, data_frame in data.items():
@@ -42,14 +43,15 @@ def write_points(self, data, *args, **kwargs):
                     end_index = (batch + 1) * batch_size
                     data = [self._convert_dataframe_to_json(
                         name=key,
-                        dataframe=data_frame.ix[start_index:end_index].copy())]
+                        dataframe=data_frame.ix[start_index:end_index].copy(),
+                        time_precision=time_precision)]
                     InfluxDBClient.write_points_with_precision(self, data,
                                                                *args, **kwargs)
             return True
         else:
-            data = [self._convert_dataframe_to_json(name=key,
-                                                    dataframe=dataframe)
-                    for key, dataframe in data.items()]
+            data = [self._convert_dataframe_to_json(
+                name=key, dataframe=dataframe, time_precision=time_precision)
+                for key, dataframe in data.items()]
             return InfluxDBClient.write_points_with_precision(self, data,
                                                               *args, **kwargs)
 
@@ -105,7 +107,7 @@ def _to_dataframe(self, json_result, time_precision):
         del dataframe['time']
         return dataframe
 
-    def _convert_dataframe_to_json(self, dataframe, name):
+    def _convert_dataframe_to_json(self, dataframe, name, time_precision='s'):
         try:
             import pandas as pd
         except ImportError:
@@ -120,12 +122,18 @@ def _convert_dataframe_to_json(self, dataframe, name):
         dataframe.index = dataframe.index.to_datetime()
         if dataframe.index.tzinfo is None:
             dataframe.index = dataframe.index.tz_localize('UTC')
-        dataframe['time'] = [self._datetime_to_epoch(dt)
+        dataframe['time'] = [self._datetime_to_epoch(dt, time_precision)
                              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()
+    def _datetime_to_epoch(self, datetime, time_precision='s'):
+        seconds = (datetime - DataFrameClient.EPOCH).total_seconds()
+        if time_precision == 's':
+            return seconds
+        elif time_precision == 'm':
+            return seconds * 1000
+        elif time_precision == 'u':
+            return seconds * 1000000
diff --git a/tests/influxdb/misc_test.py b/tests/influxdb/misc_test.py
@@ -8,6 +8,7 @@
 from nose.tools import raises
 from datetime import timedelta
 from tests import skipIfPYpy, using_pypy
+import copy
 
 if not using_pypy:
     import pandas as pd
@@ -111,6 +112,44 @@ def test_write_points_from_dataframe_with_period_index(self):
 
             self.assertListEqual(json.loads(m.last_request.body), points)
 
+    def test_write_points_from_dataframe_with_time_precision(self):
+        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, 0],
+                    ["2", 2, 2.0, 3600]
+                ],
+                "name": "foo",
+                "columns": ["column_one", "column_two", "column_three", "time"]
+            }
+        ]
+
+        points_ms = copy.deepcopy(points)
+        points_ms[0]["points"][1][-1] = 3600 * 1000
+
+        points_us = copy.deepcopy(points)
+        points_us[0]["points"][1][-1] = 3600 * 1000000
+
+        with requests_mock.Mocker() as m:
+            m.register_uri(requests_mock.POST,
+                           "http://localhost:8086/db/db/series")
+
+            cli = DataFrameClient(database='db')
+
+            cli.write_points({"foo": dataframe}, time_precision='s')
+            self.assertListEqual(json.loads(m.last_request.body), points)
+
+            cli.write_points({"foo": dataframe}, time_precision='m')
+            self.assertListEqual(json.loads(m.last_request.body), points_ms)
+
+            cli.write_points({"foo": dataframe}, time_precision='u')
+            self.assertListEqual(json.loads(m.last_request.body), points_us)
+
     @raises(TypeError)
     def test_write_points_from_dataframe_fails_without_time_index(self):
         dataframe = pd.DataFrame(data=[["1", 1, 1.0], ["2", 2, 2.0]],
