made a separate FutureDayBarStore out of DayBarStore

rqalpha/data/base_data_source/data_source.py

@@ -37,9 +37,10 @@
                                 AbstractDayBarStore, AbstractDividendStore,
                                 AbstractInstrumentStore)
 from .storages import (DateSet, DayBarStore, DividendStore,
-                       ExchangeTradingCalendarStore, FutureInfoStore,
-                       InstrumentStore, ShareTransformationStore,
-                       SimpleFactorStore, YieldCurveStore)
+                       ExchangeTradingCalendarStore, FutureDayBarStore,
+                       FutureInfoStore, InstrumentStore,
+                       ShareTransformationStore, SimpleFactorStore,
+                       YieldCurveStore)
 
 
 class BaseDataSource(AbstractDataSource):
@@ -59,7 +60,7 @@ def _p(name):
         self._day_bars = {
             INSTRUMENT_TYPE.CS: DayBarStore(_p('stocks.h5')),
             INSTRUMENT_TYPE.INDX: DayBarStore(_p('indexes.h5')),
-            INSTRUMENT_TYPE.FUTURE: DayBarStore(_p('futures.h5')),
+            INSTRUMENT_TYPE.FUTURE: FutureDayBarStore(_p('futures.h5')),
             INSTRUMENT_TYPE.ETF: funds_day_bar_store,
             INSTRUMENT_TYPE.LOF: funds_day_bar_store
         }  # type: Dict[INSTRUMENT_TYPE, AbstractDayBarStore]
@@ -170,7 +171,7 @@ def is_st_stock(self, order_book_id, dates):
 
     @lru_cache(None)
     def _all_day_bars_of(self, instrument):
-        return self._day_bars[instrument.type].get_bars(instrument.order_book_id, instrument=instrument) 
+        return self._day_bars[instrument.type].get_bars(instrument.order_book_id) 
 
     @lru_cache(None)
     def _filtered_day_bars(self, instrument):

rqalpha/data/base_data_source/storages.py

@@ -168,22 +168,11 @@ def __init__(self, path):
             raise FileExistsError("File {} not exist，please update bundle.".format(path))
         self._h5 = open_h5(path, mode="r")
 
-    def get_bars(self, order_book_id, instrument=None):
+    def get_bars(self, order_book_id):
         try:
-            bars = self._h5[order_book_id][:]
+            return self._h5[order_book_id][:]
         except KeyError:
-            bars = np.empty(0, dtype=self.DEFAULT_DTYPE)
-        if instrument is not None and instrument.type == 'Future':
-            new_fields = [field for field in FUTURES_MISSING_FIELDS if field not in bars.dtype.names]
-            if new_fields:
-                new_dt = np.dtype(bars.dtype.descr + [(field, '<f8') for field in new_fields])
-                b = np.zeros(bars.shape, dtype=new_dt)
-                for name in bars.dtype.names:
-                    b[name] = bars[name]
-                for field in new_fields:
-                    b[field] = np.nan
-                bars = b
-        return bars
+            return np.empty(0, dtype=self.DEFAULT_DTYPE)
 
     def get_date_range(self, order_book_id):
         try:
@@ -192,6 +181,26 @@ def get_date_range(self, order_book_id):
         except KeyError:
             return 20050104, 20050104
 
+class FutureDayBarStore(DayBarStore):
+    def __init__(self, *args, **kwargs):
+        self.DEFAULT_DTYPE = np.dtype(self.DEFAULT_DTYPE.descr + [(field, '<f8') for field in FUTURES_MISSING_FIELDS])
+        super().__init__(*args, **kwargs)
+
+    def get_bars(self, order_book_id):
+        try:
+            bars = self._h5[order_book_id][:]
+        except KeyError:
+            bars = np.empty(0, dtype=self.DEFAULT_DTYPE)
+        new_fields = [field for field in FUTURES_MISSING_FIELDS if field not in bars.dtype.names]
+        if new_fields:
+            new_dt = np.dtype(bars.dtype.descr + [(field, '<f8') for field in new_fields])
+            b = np.zeros(bars.shape, dtype=new_dt)
+            for name in bars.dtype.names:
+                b[name] = bars[name]
+            for field in new_fields:
+                b[field] = np.nan
+            bars = b
+        return bars
 
 class DividendStore(AbstractDividendStore):
     def __init__(self, path):