Fix tz localize (influxdata#684)

clslgrnc · xginn8 · commit f03f4957c70a · 2019-03-14T10:04:32.000-04:00
* fix already tz-aware error

* fix tests tz_localize

* update CHANGELOG.md
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](http://semver.org/spec/v2.0.0.
 ### Added
 
 ### Changed
+- Fix 'TypeError: Already tz-aware' introduced with recent versions of Panda (#671, #676, thx @f4bsch @clslgrnc)
 
 ### Removed
 
diff --git a/influxdb/_dataframe_client.py b/influxdb/_dataframe_client.py
@@ -202,7 +202,8 @@ def _to_dataframe(self, rs, dropna=True):
             df = pd.DataFrame(data)
             df.time = pd.to_datetime(df.time)
             df.set_index('time', inplace=True)
-            df.index = df.index.tz_localize('UTC')
+            if df.index.tzinfo is None:
+                df.index = df.index.tz_localize('UTC')
             df.index.name = None
             result[key].append(df)
         for key, data in result.items():
diff --git a/influxdb/tests/dataframe_client_test.py b/influxdb/tests/dataframe_client_test.py
@@ -818,13 +818,15 @@ def test_query_into_dataframe(self):
         pd1 = pd.DataFrame(
             [[23422]], columns=['value'],
             index=pd.to_datetime(["2009-11-10T23:00:00Z"]))
-        pd1.index = pd1.index.tz_localize('UTC')
+        if pd1.index.tzinfo is None:
+            pd1.index = pd1.index.tz_localize('UTC')
         pd2 = pd.DataFrame(
             [[23422], [23422], [23422]], columns=['value'],
             index=pd.to_datetime(["2009-11-10T23:00:00Z",
                                   "2009-11-10T23:00:00Z",
                                   "2009-11-10T23:00:00Z"]))
-        pd2.index = pd2.index.tz_localize('UTC')
+        if pd2.index.tzinfo is None:
+            pd2.index = pd2.index.tz_localize('UTC')
         expected = {
             ('network', (('direction', ''),)): pd1,
             ('network', (('direction', 'in'),)): pd2
@@ -871,11 +873,14 @@ def test_multiquery_into_dataframe(self):
             index=pd.to_datetime([
                 "2015-01-29 21:55:43.702900257+0000",
                 "2015-01-29 21:55:43.702900257+0000",
-                "2015-06-11 20:46:02+0000"])).tz_localize('UTC')
+                "2015-06-11 20:46:02+0000"]))
+        if pd1.index.tzinfo is None:
+            pd1.index = pd1.index.tz_localize('UTC')
         pd2 = pd.DataFrame(
             [[3]], columns=['count'],
-            index=pd.to_datetime(["1970-01-01 00:00:00+00:00"]))\
-            .tz_localize('UTC')
+            index=pd.to_datetime(["1970-01-01 00:00:00+00:00"]))
+        if pd2.index.tzinfo is None:
+            pd2.index = pd2.index.tz_localize('UTC')
         expected = [{'cpu_load_short': pd1}, {'cpu_load_short': pd2}]
 
         cli = DataFrameClient('host', 8086, 'username', 'password', 'db')
