Slightly better dataframe_client.

Grégory Starck · Grégory Starck · commit 1c5fed3568ae · 2015-03-25T09:49:03.000-04:00
diff --git a/influxdb/_dataframe_client.py b/influxdb/_dataframe_client.py
@@ -7,10 +7,7 @@
 
 from .client import InfluxDBClient
 
-try:
-    import pandas as pd
-except ImportError:
-    pd = None
+import pandas as pd
 
 
 class DataFrameClient(InfluxDBClient):
@@ -20,14 +17,7 @@ class DataFrameClient(InfluxDBClient):
     The client reads and writes from pandas DataFrames.
     """
 
-    def __init__(self, *args, **kwargs):
-        super(DataFrameClient, self).__init__(*args, **kwargs)
-        if not pd:
-            raise ImportError(
-                'DataFrameClient requires Pandas'
-            )
-
-        self.EPOCH = pd.Timestamp('1970-01-01 00:00:00.000+00:00')
+    EPOCH = pd.Timestamp('1970-01-01 00:00:00.000+00:00')
 
     def write_points(self, data, *args, **kwargs):
         """
@@ -56,13 +46,15 @@ def write_points(self, data, *args, **kwargs):
                         name=key,
                         dataframe=data_frame.ix[start_index:end_index].copy(),
                         time_precision=time_precision)]
-                    InfluxDBClient.write_points(self, data, *args, **kwargs)
+                    super(DataFrameClient, self).write_points(data,
+                                                              *args, **kwargs)
             return True
         else:
             data = [self._convert_dataframe_to_json(
                 name=key, dataframe=dataframe, time_precision=time_precision)
                 for key, dataframe in data.items()]
-            return InfluxDBClient.write_points(self, data, *args, **kwargs)
+            return super(DataFrameClient, self).write_points(data,
+                                                             *args, **kwargs)
 
     def write_points_with_precision(self, data, time_precision='s'):
         """
@@ -86,14 +78,11 @@ def query(self, query, time_precision='s', chunked=False):
             retrieved in chunks, False otherwise.
 
         """
-        result = super(DataFrameClient, self).query(
-                                      query=query,
-                                      time_precision=time_precision,
-                                      chunked=chunked)
-        if len(result['results'][0]) > 0:
-            return self._to_dataframe(result['results'][0], time_precision)
+        results = super(DataFrameClient, self).query(query, database=database)
+        if len(results) > 0:
+            return self._to_dataframe(results, time_precision)
         else:
-            return result
+            return results
 
     def _to_dataframe(self, json_result, time_precision):
         dataframe = pd.DataFrame(data=json_result['points'],
diff --git a/influxdb/dataframe_client.py b/influxdb/dataframe_client.py
@@ -0,0 +1,19 @@
+# -*- coding: utf-8 -*-
+"""
+DataFrame client for InfluxDB
+"""
+
+__all__ = ['DataFrameClient']
+
+try:
+    import pandas
+    del pandas
+except ImportError as err:
+    from .client import InfluxDBClient
+
+    class DataFrameClient(InfluxDBClient):
+        def __init__(self, *a, **kw):
+            raise ImportError("DataFrameClient requires Pandas "
+                              "which couldn't be imported: %s" % err)
+else:
+    from ._dataframe_client import DataFrameClient
diff --git a/tests/influxdb/dataframe_client_test.py b/tests/influxdb/dataframe_client_test.py
@@ -211,13 +211,24 @@ def test_query_with_empty_result(self):
             assert result == []
 
     def test_list_series(self):
-        response = [
-            {
-                'columns': ['time', 'name'],
-                'name': 'list_series_result',
-                'points': [[0, 'seriesA'], [0, 'seriesB']]
-            }
-        ]
+        response = {
+            'results': [
+                {
+                    'series': [{
+                        'columns': ['id'],
+                        'name': 'seriesA',
+                        'values': [[0]],
+                    }]
+                },
+                {
+                    'series': [{
+                        'columns': ['id'],
+                        'name': 'seriesB',
+                        'values': [[1]],
+                    }]
+                },
+            ]
+        }
         with _mocked_session('get', 200, response):
             cli = DataFrameClient('host', 8086, 'username', 'password', 'db')
             series_list = cli.get_list_series()
