Give dates correct tzinfo after mongo query (pandas-dev#556)

Author: srf94

arctic/tickstore/toplevel.py

@@ -152,7 +152,14 @@ def _get_libraries(self, date_range):
 
         rtn = [TickStoreLibrary(self._arctic_lib.arctic[library.library], library.date_range)
                for library in libraries]
-        current_start = rtn[-1].date_range.end if rtn else dt(1970, 1, 1, 0, 0)  # epoch
+
+        if rtn:
+            current_start = rtn[-1].date_range.end
+        elif date_range.end.tzinfo:
+            current_start = dt(1970, 1, 1, 0, 0, tzinfo=date_range.end.tzinfo)
+        else:
+            current_start = dt(1970, 1, 1, 0, 0)
+
         if date_range.end is None or current_start < date_range.end:
             name = self.get_name()
             db_name, tick_type = name.split('.', 1)
@@ -195,8 +202,20 @@ def _get_library_metadata(self, date_range):
         query = {'$or': [{'start': {'$lte': start}, 'end': {'$gte': start}},
                          {'start': {'$gte': start}, 'end': {'$lte': end}},
                          {'start': {'$lte': end}, 'end': {'$gte': end}}]}
-        return [TickStoreLibrary(res['library_name'], DateRange(res['start'], res['end'], CLOSED_CLOSED))
-                for res in self._collection.find(query,
-                                                 projection={'library_name': 1,
-                                                             'start': 1, 'end': 1},
-                                                 sort=[('start', pymongo.ASCENDING)])]
+
+        cursor = self._collection.find(query,
+                                       projection={'library_name': 1, 'start': 1, 'end': 1},
+                                       sort=[('start', pymongo.ASCENDING)])
+
+        results = []
+        for res in cursor:
+            start = res['start']
+            if date_range.start.tzinfo is not None and start.tzinfo is None:
+                start = start.replace(tzinfo=mktz("UTC")).astimezone(tz=date_range.start.tzinfo)
+
+            end = res['end']
+            if date_range.end.tzinfo is not None and end.tzinfo is None:
+                end = end.replace(tzinfo=mktz("UTC")).astimezone(tz=date_range.end.tzinfo)
+
+            results.append(TickStoreLibrary(res['library_name'], DateRange(start, end, CLOSED_CLOSED)))
+        return results

tests/integration/tickstore/test_toplevel.py

@@ -242,3 +242,31 @@ def test_no_min_max_date(arctic):
         tstore.min_date('unknown-symbol')
     with pytest.raises(NoDataFoundException):
         tstore.max_date('unknown-symbol')
+
+
+def test_get_libraries_no_data_raises_exception(toplevel_tickstore, arctic):
+    date_range = DateRange(start=dt(2009, 1, 1), end=dt(2010, 12, 31, 23, 59, 59, 999000))
+    with pytest.raises(NoDataFoundException):
+        toplevel_tickstore._get_libraries(date_range)
+
+
+def test_get_libraries_no_data_raises_exception_tzinfo_given(toplevel_tickstore, arctic):
+    tzinfo = mktz('Asia/Chongqing')
+    date_range = DateRange(start=dt(2009, 1, 1, tzinfo=tzinfo),
+                           end=dt(2010, 12, 31, 23, 59, 59, 999000, tzinfo=tzinfo))
+    with pytest.raises(NoDataFoundException):
+        toplevel_tickstore._get_libraries(date_range)
+
+
+def test_get_library_metadata(arctic):
+    arctic.initialize_library('FEED_2010.LEVEL1', tickstore.TICK_STORE_TYPE)
+    arctic.initialize_library('FEED_2011.LEVEL1', tickstore.TICK_STORE_TYPE)
+    arctic.initialize_library('FEED.LEVEL1', toplevel.TICK_STORE_TYPE)
+    toplevel_tickstore = arctic['FEED.LEVEL1']
+
+    symbol = "USD"
+    tzinfo=mktz('Asia/Chongqing')
+    with pytest.raises(NoDataFoundException):
+        toplevel_tickstore.read(symbol, DateRange(start=dt(2010, 1, 1, tzinfo=tzinfo),
+                                                  end=dt(2011, 1, 2, tzinfo=tzinfo)),
+                                columns=None)