Minor changes to _stringify_dataframe. Added test for numeric precision.

mdbartos · mdbartos · commit 56062c5160bc · 2016-08-30T03:00:45.000-04:00
diff --git a/influxdb/_dataframe_client.py b/influxdb/_dataframe_client.py
@@ -305,7 +305,7 @@ def _stringify_dataframe(self,
                              datatype='field'):
 
         # Find int and string columns for field-type data
-        int_columns = dataframe.select_dtypes(include=['int']).columns
+        int_columns = dataframe.select_dtypes(include=['integer']).columns
         string_columns = dataframe.select_dtypes(include=['object']).columns
 
         # Convert dataframe to string
@@ -323,11 +323,20 @@ def _stringify_dataframe(self,
                                            .astype(str))
         elif isinstance(numeric_precision, int):
             # If precision is specified, round to appropriate precision
-            numeric_columns = (dataframe.select_dtypes(include=['number'])
-                               .columns)
-            dataframe[numeric_columns] = (dataframe[numeric_columns]
-                                          .round(numeric_precision))
-            dataframe = dataframe.astype(str)
+            float_columns = (dataframe.select_dtypes(include=['floating'])
+                             .columns)
+            nonfloat_columns = dataframe.columns[~dataframe.columns.isin(
+                float_columns)]
+            dataframe[float_columns] = (dataframe[float_columns]
+                                        .round(numeric_precision))
+            # If desired precision is > 10 decimal places, need to use repr
+            if numeric_precision > 10:
+                dataframe[float_columns] = (dataframe[float_columns]
+                                            .applymap(repr))
+                dataframe[nonfloat_columns] = (dataframe[nonfloat_columns]
+                                               .astype(str))
+            else:
+                dataframe = dataframe.astype(str)
         else:
             raise ValueError('Invalid numeric precision')
 
diff --git a/influxdb/tests/dataframe_client_test.py b/influxdb/tests/dataframe_client_test.py
@@ -69,7 +69,7 @@ def test_write_points_from_dataframe_in_batches(self):
             cli = DataFrameClient(database='db')
             self.assertTrue(cli.write_points(dataframe, "foo", batch_size=1))
 
-    def test_write_points_with_tag_columns(self):
+    def test_write_points_from_dataframe_with_tag_columns(self):
         now = pd.Timestamp('1970-01-01 00:00+00:00')
         dataframe = pd.DataFrame(data=[['blue', 1, "1", 1, 1.0],
                                        ['red', 0, "2", 2, 2.0]],
@@ -100,7 +100,7 @@ def test_write_points_with_tag_columns(self):
                              tag_columns=['tag_one', 'tag_two'], tags=None)
             self.assertEqual(m.last_request.body, expected)
 
-    def test_write_points_with_tag_columns_and_global_tags(self):
+    def test_write_points_from_dataframe_with_tag_cols_and_global_tags(self):
         now = pd.Timestamp('1970-01-01 00:00+00:00')
         dataframe = pd.DataFrame(data=[['blue', 1, "1", 1, 1.0],
                                        ['red', 0, "2", 2, 2.0]],
@@ -128,7 +128,7 @@ def test_write_points_with_tag_columns_and_global_tags(self):
                              tags={'global_tag': 'value'})
             self.assertEqual(m.last_request.body, expected)
 
-    def test_write_points_with_tag_columns_and_defaults(self):
+    def test_write_points_from_dataframe_with_tag_cols_and_defaults(self):
         now = pd.Timestamp('1970-01-01 00:00+00:00')
         dataframe = pd.DataFrame(data=[['blue', 1, "1", 1, 1.0, 'hot'],
                                        ['red', 0, "2", 2, 2.0, 'cold']],
@@ -219,6 +219,50 @@ def test_write_points_from_dataframe_with_numeric_column_names(self):
 
             self.assertEqual(m.last_request.body, expected)
 
+    def test_write_points_from_dataframe_with_numeric_precision(self):
+        now = pd.Timestamp('1970-01-01 00:00+00:00')
+        # df with numeric column names
+        dataframe = pd.DataFrame(data=[["1", 1, 1.1111111111111],
+                                       ["2", 2, 2.2222222222222]],
+                                 index=[now, now + timedelta(hours=1)])
+
+        expected_default_precision = (
+            b'foo,hello=there 0=\"1\",1=1i,2=1.11111111111 0\n'
+            b'foo,hello=there 0=\"2\",1=2i,2=2.22222222222 3600000000000\n'
+        )
+
+        expected_specified_precision = (
+            b'foo,hello=there 0=\"1\",1=1i,2=1.1111 0\n'
+            b'foo,hello=there 0=\"2\",1=2i,2=2.2222 3600000000000\n'
+        )
+
+        expected_full_precision = (
+            b'foo,hello=there 0=\"1\",1=1i,2=1.1111111111111 0\n'
+            b'foo,hello=there 0=\"2\",1=2i,2=2.2222222222222 3600000000000\n'
+        )
+
+        with requests_mock.Mocker() as m:
+            m.register_uri(requests_mock.POST,
+                           "http://localhost:8086/write",
+                           status_code=204)
+
+            cli = DataFrameClient(database='db')
+            cli.write_points(dataframe, "foo", {"hello": "there"})
+
+            self.assertEqual(m.last_request.body, expected_default_precision)
+
+            cli = DataFrameClient(database='db')
+            cli.write_points(dataframe, "foo", {"hello": "there"},
+                             numeric_precision=4)
+
+            self.assertEqual(m.last_request.body, expected_specified_precision)
+
+            cli = DataFrameClient(database='db')
+            cli.write_points(dataframe, "foo", {"hello": "there"},
+                             numeric_precision='full')
+
+            self.assertEqual(m.last_request.body, expected_full_precision)
+
     def test_write_points_from_dataframe_with_period_index(self):
         dataframe = pd.DataFrame(data=[["1", 1, 1.0], ["2", 2, 2.0]],
                                  index=[pd.Period('1970-01-01'),
