Merge pull request pandas-dev#97 from manahl/top_level_read_fix

Top level read fix for pandas-dev#95
dan-nadler · Jan 28, 2016 · 7e14828 · 7e14828
2 parents db7adb8 + c115ee9
commit 7e14828
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 8 deletions.
diff --git a/arctic/tickstore/toplevel.py b/arctic/tickstore/toplevel.py
@@ -100,8 +100,16 @@ def add(self, date_range, library_name):
 
     def read(self, symbol, date_range, columns=['BID', 'ASK', 'TRDPRC_1', 'BIDSIZE', 'ASKSIZE', 'TRDVOL_1'], include_images=False):
         libraries = self._get_libraries(date_range)
-        dfs = [l.library.read(symbol, l.date_range.intersection(date_range), columns,
-                              include_images=include_images) for l in libraries]
+        dfs = []
+        for l in libraries:
+            try:
+                df = l.library.read(symbol, l.date_range.intersection(date_range), columns,
+                                    include_images=include_images)
+                dfs.append(df)
+            except NoDataFoundException as e:
+                continue
+        if len(dfs) == 0:
+            raise NoDataFoundException("No Data found for {} in range: {}".format(symbol, date_range))
         return pd.concat(dfs)
 
     def write(self, symbol, data):

diff --git a/tests/integration/tickstore/test_toplevel.py b/tests/integration/tickstore/test_toplevel.py
@@ -71,12 +71,8 @@ def test_should_return_data_when_date_range_spans_libraries(toplevel_tickstore,
     arctic.initialize_library('FEED_2011.LEVEL1', tickstore.TICK_STORE_TYPE)
     tickstore_2010 = arctic['FEED_2010.LEVEL1']
     tickstore_2011 = arctic['FEED_2011.LEVEL1']
-    toplevel_tickstore._collection.insert_one({'start': dt(2010, 1, 1),
-                                           'end': dt(2010, 12, 31, 23, 59, 59),
-                                           'library_name': 'FEED_2010.LEVEL1'})
-    toplevel_tickstore._collection.insert_one({'start': dt(2011, 1, 1),
-                                           'end': dt(2011, 12, 31, 23, 59, 59),
-                                           'library_name': 'FEED_2011.LEVEL1'})
+    toplevel_tickstore.add(DateRange(start=dt(2010, 1, 1), end=dt(2010, 12, 31, 23, 59, 59, 999000)), 'FEED_2010.LEVEL1')
+    toplevel_tickstore.add(DateRange(start=dt(2011, 1, 1), end=dt(2011, 12, 31, 23, 59, 59, 999000)), 'FEED_2011.LEVEL1')
     dates = pd.date_range('20100101', periods=6, tz=mktz('Europe/London'))
     df_10 = pd.DataFrame(np.random.randn(6, 4), index=dates, columns=list('ABCD'))
     tickstore_2010.write('blah', df_10)
@@ -88,6 +84,24 @@ def test_should_return_data_when_date_range_spans_libraries(toplevel_tickstore,
     assert_frame_equal(expected_df, res.tz_convert(mktz('Europe/London')))
 
 
+def test_should_return_data_when_date_range_spans_libraries_even_if_one_returns_nothing(toplevel_tickstore, arctic):
+    arctic.initialize_library('FEED_2010.LEVEL1', tickstore.TICK_STORE_TYPE)
+    arctic.initialize_library('FEED_2011.LEVEL1', tickstore.TICK_STORE_TYPE)
+    tickstore_2010 = arctic['FEED_2010.LEVEL1']
+    tickstore_2011 = arctic['FEED_2011.LEVEL1']
+    toplevel_tickstore.add(DateRange(start=dt(2010, 1, 1), end=dt(2010, 12, 31, 23, 59, 59, 999000)), 'FEED_2010.LEVEL1')
+    toplevel_tickstore.add(DateRange(start=dt(2011, 1, 1), end=dt(2011, 12, 31, 23, 59, 59, 999000)), 'FEED_2011.LEVEL1')
+    dates = pd.date_range('20100101', periods=6, tz=mktz('Europe/London'))
+    df_10 = pd.DataFrame(np.random.randn(6, 4), index=dates, columns=list('ABCD'))
+    tickstore_2010.write('blah', df_10)
+    dates = pd.date_range('20110201', periods=6, tz=mktz('Europe/London'))
+    df_11 = pd.DataFrame(np.random.randn(6, 4), index=dates, columns=list('ABCD'))
+    tickstore_2011.write('blah', df_11)
+    res = toplevel_tickstore.read('blah', DateRange(start=dt(2010, 1, 2), end=dt(2011, 1, 4)), list('ABCD'))
+    expected_df = df_10[1:]
+    assert_frame_equal(expected_df, res.tz_convert(mktz('Europe/London')))
+
+
 def test_should_add_underlying_library_where_none_exists(toplevel_tickstore, arctic):
     arctic.initialize_library('FEED_2010.LEVEL1', tickstore.TICK_STORE_TYPE)
     toplevel_tickstore.add(DateRange(start=dt(2010, 1, 1), end=dt(2010, 12, 31, 23, 59, 59, 999000)), 'FEED_2010.LEVEL1')