Modify data.storage

qlib/data/data.py

@@ -6,6 +6,7 @@ from __future__ import division
 from __future__ import print_function
 
 import os
+import re
 import abc
 import time
 import queue
@@ -27,12 +28,35 @@ from .cache import DiskDatasetCache, DiskExpressionCache
 from ..utils import Wrapper, init_instance_by_config, register_wrapper, get_module_by_module_path
 
 
-class CalendarProvider(abc.ABC):
+class ProviderBackendMixin:
+    def get_default_backend(self):
+        backend = {}
+        provider_name = re.findall("[A-Z][^A-Z]*", self.__class__.__name__)[-2]  # type: str
+        # set default storage class
+        backend.setdefault("class", f"File{provider_name}Storage")
+        # set default storage module
+        backend.setdefault("module_path", "qlib.data.storage.file_storage")
+        # set default storage kwargs
+        backend_kwargs = backend.setdefault("kwargs", {})  # type: dict
+        backend_kwargs.setdefault("uri", os.path.join(C.get_data_path(), f"{provider_name.lower()}s"))
+        return backend
+
+    @property
+    def backend_obj(self):
+        return init_instance_by_config(self.backend)
+
+
+class CalendarProvider(abc.ABC, ProviderBackendMixin):
     """Calendar provider base class
 
     Provide calendar data.
     """
 
+    def __init__(self, *args, **kwargs):
+        self.backend = kwargs.get("backend", {})
+        if not self.backend:
+            self.backend = self.get_default_backend()
+
     @abc.abstractmethod
     def calendar(self, start_time=None, end_time=None, freq="day", future=False):
         """Get calendar of certain market in given time range.
@@ -127,12 +151,17 @@ class CalendarProvider(abc.ABC):
         return hash_args(start_time, end_time, freq, future)
 
 
-class InstrumentProvider(abc.ABC):
+class InstrumentProvider(abc.ABC, ProviderBackendMixin):
     """Instrument provider base class
 
     Provide instrument data.
     """
 
+    def __init__(self, *args, **kwargs):
+        self.backend = kwargs.get("backend", {})
+        if not self.backend:
+            self.backend = self.get_default_backend()
+
     @staticmethod
     def instruments(market="all", filter_pipe=None):
         """Get the general config dictionary for a base market adding several dynamic filters.
@@ -215,12 +244,17 @@ class InstrumentProvider(abc.ABC):
         raise ValueError(f"Unknown instrument type {inst}")
 
 
-class FeatureProvider(abc.ABC):
+class FeatureProvider(abc.ABC, ProviderBackendMixin):
     """Feature provider class
 
     Provide feature data.
     """
 
+    def __init__(self, *args, **kwargs):
+        self.backend = kwargs.get("backend", {})
+        if not self.backend:
+            self.backend = self.get_default_backend()
+
     @abc.abstractmethod
     def feature(self, instrument, field, start_time, end_time, freq):
         """Get feature data.
@@ -497,6 +531,7 @@ class LocalCalendarProvider(CalendarProvider):
     """
 
     def __init__(self, **kwargs):
+        super(LocalCalendarProvider, self).__init__(**kwargs)
         self.remote = kwargs.get("remote", False)
 
     @property
@@ -517,18 +552,8 @@ class LocalCalendarProvider(CalendarProvider):
         list
             list of timestamps
         """
-        if future:
-            fname = self._uri_cal.format(freq + "_future")
-            # if future calendar not exists, return current calendar
-            if not os.path.exists(fname):
-                get_module_logger("data").warning(f"{freq}_future.txt not exists, return current calendar!")
-                fname = self._uri_cal.format(freq)
-        else:
-            fname = self._uri_cal.format(freq)
-        if not os.path.exists(fname):
-            raise ValueError("calendar not exists for freq " + freq)
-        with open(fname) as f:
-            return [pd.Timestamp(x.strip()) for x in f]
+        self.backend.setdefault("kwargs", {}).update(freq=freq, future=future)
+        return [pd.Timestamp(x) for x in self.backend_obj.data]
 
     def calendar(self, start_time=None, end_time=None, freq="day", future=False):
         _calendar, _calendar_index = self._get_calendar(freq, future)
@@ -559,31 +584,15 @@ class LocalInstrumentProvider(InstrumentProvider):
     Provide instrument data from local data source.
     """
 
-    def __init__(self):
-        pass
-
     @property
     def _uri_inst(self):
         """Instrument file uri."""
         return os.path.join(C.get_data_path(), "instruments", "{}.txt")
 
     def _load_instruments(self, market):
-        fname = self._uri_inst.format(market)
-        if not os.path.exists(fname):
-            raise ValueError("instruments not exists for market " + market)
 
-        _instruments = dict()
-        df = pd.read_csv(
-            fname,
-            sep="\t",
-            usecols=[0, 1, 2],
-            names=["inst", "start_datetime", "end_datetime"],
-            dtype={"inst": str},
-            parse_dates=["start_datetime", "end_datetime"],
-        )
-        for row in df.itertuples(index=False):
-            _instruments.setdefault(row[0], []).append((row[1], row[2]))
-        return _instruments
+        self.backend.setdefault("kwargs", {}).update(market=market)
+        return self.backend_obj.data
 
     def list_instruments(self, instruments, start_time=None, end_time=None, freq="day", as_list=False):
         market = instruments["market"]
@@ -601,7 +610,7 @@ class LocalInstrumentProvider(InstrumentProvider):
             inst: list(
                 filter(
                     lambda x: x[0] <= x[1],
-                    [(max(start_time, x[0]), min(end_time, x[1])) for x in spans],
+                    [(max(start_time, pd.Timestamp(x[0])), min(end_time, pd.Timestamp(x[1]))) for x in spans],
                 )
             )
             for inst, spans in _instruments.items()
@@ -627,6 +636,7 @@ class LocalFeatureProvider(FeatureProvider):
     """
 
     def __init__(self, **kwargs):
+        super(LocalFeatureProvider, self).__init__(**kwargs)
         self.remote = kwargs.get("remote", False)
 
     @property
@@ -638,14 +648,9 @@ class LocalFeatureProvider(FeatureProvider):
         # validate
         field = str(field).lower()[1:]
         instrument = code_to_fname(instrument)
-        uri_data = self._uri_data.format(instrument.lower(), field, freq)
-        if not os.path.exists(uri_data):
-            get_module_logger("data").warning("WARN: data not found for %s.%s" % (instrument, field))
-            return pd.Series(dtype=np.float32)
-            # raise ValueError('uri_data not found: ' + uri_data)
-        # load
-        series = read_bin(uri_data, start_index, end_index)
-        return series
+
+        self.backend.setdefault("kwargs", {}).update(instrument=instrument, field=field, freq=freq)
+        return self.backend_obj[start_index : end_index + 1]
 
 
 class LocalExpressionProvider(ExpressionProvider):
@@ -1061,7 +1066,8 @@ def register_all_wrappers(C):
     register_wrapper(Cal, _calendar_provider, "qlib.data")
     logger.debug(f"registering Cal {C.calendar_provider}-{C.calendar_cache}")
 
-    register_wrapper(Inst, C.instrument_provider, "qlib.data")
+    _instrument_provider = init_instance_by_config(C.instrument_provider, module)
+    register_wrapper(Inst, _instrument_provider, "qlib.data")
     logger.debug(f"registering Inst {C.instrument_provider}")
 
     if getattr(C, "feature_provider", None) is not None:

qlib/data/storage/file_storage.py

@@ -8,24 +8,29 @@ from typing import Iterator, Iterable, Union, Dict, Mapping, Tuple
 import numpy as np
 import pandas as pd
 
-from . import CalendarStorage, InstrumentStorage, FeatureStorage, CalVT, InstKT, InstVT
+from qlib.data.storage import CalendarStorage, InstrumentStorage, FeatureStorage, CalVT, InstKT, InstVT
 
 
 class FileCalendarStorage(CalendarStorage):
-    def __init__(self, uri: str):
-        super(FileCalendarStorage, self).__init__(uri)
-        self._uri = Path(self._uri).expanduser().resolve()
+    def __init__(self, freq: str, future: bool, uri: str):
+        super(FileCalendarStorage, self).__init__(freq, future, uri)
+        _file_name = f"{freq}_future.txt" if future else f"{freq}.txt"
+        self.uri = Path(self.uri).expanduser().joinpath(_file_name.lower())
 
     def _read_calendar(self, skip_rows: int = 0, n_rows: int = None) -> np.ndarray:
-        if not self._uri.exists():
+        if not self.uri.exists():
             self._write_calendar(values=[])
-        with self._uri.open("rb") as fp:
+        with self.uri.open("rb") as fp:
             return np.loadtxt(fp, str, skiprows=skip_rows, max_rows=n_rows, encoding="utf-8")
 
     def _write_calendar(self, values: Iterable[CalVT], mode: str = "wb"):
-        with self._uri.open(mode=mode) as fp:
+        with self.uri.open(mode=mode) as fp:
             np.savetxt(fp, values, fmt="%s", encoding="utf-8")
 
+    @property
+    def data(self) -> Iterable[CalVT]:
+        return self._read_calendar()
+
     def extend(self, values: Iterable[CalVT]) -> None:
         self._write_calendar(values, mode="ab")
 
@@ -64,27 +69,27 @@ class FileCalendarStorage(CalendarStorage):
         return len(self._read_calendar())
 
     def __iter__(self):
-        with self._uri.open("r") as fp:
-            yield fp.readline()
+        return iter(self._read_calendar())
 
 
 class FileInstrumentStorage(InstrumentStorage):
+
     INSTRUMENT_SEP = "\t"
     INSTRUMENT_START_FIELD = "start_datetime"
     INSTRUMENT_END_FIELD = "end_datetime"
     SYMBOL_FIELD_NAME = "instrument"
 
-    def __init__(self, uri: str):
-        super(FileInstrumentStorage, self).__init__(uri=uri)
-        self._uri = Path(self._uri).expanduser().resolve()
+    def __init__(self, market: str, uri: str):
+        super(FileInstrumentStorage, self).__init__(market, uri)
+        self.uri = Path(self.uri).expanduser().joinpath(f"{market.lower()}.txt")
 
     def _read_instrument(self) -> Dict[InstKT, InstVT]:
-        if not self._uri.exists():
+        if not self.uri.exists():
             self._write_instrument()
 
         _instruments = dict()
         df = pd.read_csv(
-            self._uri,
+            self.uri,
             sep="\t",
             usecols=[0, 1, 2],
             names=[self.SYMBOL_FIELD_NAME, self.INSTRUMENT_START_FIELD, self.INSTRUMENT_END_FIELD],
@@ -97,7 +102,7 @@ class FileInstrumentStorage(InstrumentStorage):
 
     def _write_instrument(self, data: Dict[InstKT, InstVT] = None) -> None:
         if not data:
-            with self._uri.open("w") as _:
+            with self.uri.open("w") as _:
                 pass
             return
 
@@ -109,13 +114,17 @@ class FileInstrumentStorage(InstrumentStorage):
 
         df = pd.concat(res, sort=False)
         df.loc[:, [self.SYMBOL_FIELD_NAME, self.INSTRUMENT_START_FIELD, self.INSTRUMENT_END_FIELD]].to_csv(
-            self._uri, header=False, sep=self.INSTRUMENT_SEP, index=False
+            self.uri, header=False, sep=self.INSTRUMENT_SEP, index=False
         )
-        df.to_csv(self._uri, sep="\t", encoding="utf-8", header=False, index=False)
+        df.to_csv(self.uri, sep="\t", encoding="utf-8", header=False, index=False)
 
     def clear(self) -> None:
         self._write_instrument(data={})
 
+    @property
+    def data(self) -> Dict[InstKT, InstVT]:
+        return self._read_instrument()
+
     def __setitem__(self, k: InstKT, v: InstVT) -> None:
         inst = self._read_instrument()
         inst[k] = v
@@ -143,7 +152,7 @@ class FileInstrumentStorage(InstrumentStorage):
             raise TypeError(f"update expected at most 1 arguments, got {len(args)}")
         inst = self._read_instrument()
         if args:
-            other = args[0]
+            other = args[0]  # type: dict
             if isinstance(other, Mapping):
                 for key in other:
                     inst[key] = other[key]
@@ -160,29 +169,35 @@ class FileInstrumentStorage(InstrumentStorage):
 
 
 class FileFeatureStorage(FeatureStorage):
-    def __init__(self, uri: str):
-        super(FileFeatureStorage, self).__init__(uri=uri)
-        self._uri = Path(self._uri)
+    def __init__(self, instrument: str, field: str, freq: str, uri: str):
+        super(FileFeatureStorage, self).__init__(instrument, field, freq, uri)
+        self.uri = (
+            Path(self.uri).expanduser().joinpath(instrument.lower()).joinpath(f"{field.lower()}.{freq.lower()}.bin")
+        )
 
     def clear(self):
-        with self._uri.open("wb") as _:
+        with self.uri.open("wb") as _:
             pass
 
+    @property
+    def data(self) -> pd.Series:
+        return self[:]
+
     def extend(self, series: pd.Series) -> None:
-        extend_start_index = self[0][0] + len(self) if self._uri.exists() else series.index[0]
+        extend_start_index = self[0][0] + len(self) if self.uri.exists() else series.index[0]
         series = series.reindex(pd.RangeIndex(extend_start_index, series.index[-1] + 1))
-        with self._uri.open("ab") as fp:
+        with self.uri.open("ab") as fp:
             np.array(series.values).astype("<f").tofile(fp)
 
     def rebase(self, series: pd.Series) -> None:
         origin_series = self[:]
         series = series.append(origin_series.loc[origin_series.index > series.index[-1]])
         series = series.reindex(pd.RangeIndex(series.index[0], series.index[-1]))
-        with self._uri.open("wb") as fp:
+        with self.uri.open("wb") as fp:
             np.array(series.values).astype("<f").tofile(fp)
 
     def __getitem__(self, i: Union[int, slice]) -> Union[Tuple[int, float], pd.Series]:
-        if not self._uri.exists():
+        if not self.uri.exists():
             if isinstance(i, int):
                 return None, None
             elif isinstance(i, slice):
@@ -190,14 +205,14 @@ class FileFeatureStorage(FeatureStorage):
             else:
                 raise TypeError(f"type(i) = {type(i)}")
 
-        with open(self._uri, "rb") as fp:
+        with open(self.uri, "rb") as fp:
             ref_start_index = int(np.frombuffer(fp.read(4), dtype="<f")[0])
 
             if isinstance(i, int):
                 if ref_start_index > i:
                     raise IndexError(f"{i}: start index is {ref_start_index}")
                 fp.seek(4 * (i - ref_start_index) + 4)
-                return i, struct.unpack("f", fp.read(4))
+                return i, struct.unpack("f", fp.read(4))[0]
             elif isinstance(i, slice):
                 start_index = i.start
                 end_index = i.stop - 1
@@ -213,18 +228,18 @@ class FileFeatureStorage(FeatureStorage):
                 raise TypeError(f"type(i) = {type(i)}")
 
     def __len__(self) -> int:
-        return self._uri.stat().st_size // 4 - 1 if self._uri.exists() else 0
+        return self.uri.stat().st_size // 4 - 1 if self.uri.exists() else 0
 
     def __iter__(self):
-        if not self._uri.exists():
+        if not self.uri.exists():
             return
-        with open(self._uri, "rb") as fp:
+        with open(self.uri, "rb") as fp:
             ref_start_index = int(np.frombuffer(fp.read(4), dtype="<f")[0])
             fp.seek(4)
             while True:
                 v = fp.read(4)
                 if v:
-                    yield ref_start_index, struct.unpack("f", v)
+                    yield ref_start_index, struct.unpack("f", v)[0]
                     ref_start_index += 1
                 else:
                     break

qlib/data/storage/storage.py

@@ -1,7 +1,7 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 
-from typing import Iterable, overload, Tuple, List, Text, Iterator, Union
+from typing import Iterable, overload, Tuple, List, Text, Iterator, Union, Dict
 
 import pandas as pd
 
@@ -15,8 +15,15 @@ InstKT = Text
 
 
 class CalendarStorage:
-    def __init__(self, uri: str):
-        self._uri = uri
+    def __init__(self, freq: str, future: bool, uri: str):
+        self.freq = freq
+        self.future = future
+        self.uri = uri
+
+    @property
+    def data(self) -> Iterable[CalVT]:
+        """get all data"""
+        raise NotImplementedError("Subclass of CalendarStorage must implement `data` method")
 
     def extend(self, iterable: Iterable[CalVT]) -> None:
         raise NotImplementedError("Subclass of CalendarStorage must implement `extend` method")
@@ -82,10 +89,19 @@ class CalendarStorage:
         """x.__len__() <==> len(x)"""
         raise NotImplementedError("Subclass of CalendarStorage must implement `__len__` method")
 
+    def __iter__(self):
+        raise NotImplementedError("Subclass of CalendarStorage must implement `__iter__` method")
+
 
 class InstrumentStorage:
-    def __init__(self, uri: str):
-        self._uri = uri
+    def __init__(self, market: str, uri: str):
+        self.market = market
+        self.uri = uri
+
+    @property
+    def data(self) -> Dict[InstKT, InstVT]:
+        """get all data"""
+        raise NotImplementedError("Subclass of InstrumentStorage must implement `data` method")
 
     def clear(self) -> None:
         raise NotImplementedError("Subclass of InstrumentStorage must implement `clear` method")
@@ -120,8 +136,16 @@ class InstrumentStorage:
 
 
 class FeatureStorage:
-    def __init__(self, uri: str):
-        self._uri = uri
+    def __init__(self, instrument: str, field: str, freq: str, uri: str):
+        self.instrument = instrument
+        self.field = field
+        self.freq = freq
+        self.uri = uri
+
+    @property
+    def data(self) -> pd.Series:
+        """get all data"""
+        raise NotImplementedError("Subclass of FeatureStorage must implement `data` method")
 
     def clear(self):
         """ Remove all items from FeatureStorage. """

tests/storage_tests/test_storage.py

@@ -16,15 +16,16 @@ from qlib.data.storage.file_storage import (
     FileFeatureStorage as FeatureStorage,
 )
 
-DATA_DIR = Path(__file__).parent.joinpath("test_get_data")
+_file_name = Path(__file__).name.split(".")[0]
+DATA_DIR = Path(__file__).parent.joinpath(f"{_file_name}_data")
 QLIB_DIR = DATA_DIR.joinpath("qlib")
 QLIB_DIR.mkdir(exist_ok=True, parents=True)
 
 
 # TODO: set value
-CALENDAR_URI = QLIB_DIR.joinpath("calendars").joinpath("day.txt")
-INSTRUMENT_URI = QLIB_DIR.joinpath("instruments").joinpath("csi300.txt")
-FEATURE_URI = QLIB_DIR.joinpath("features").joinpath("SH600004").joinpath("close.day.bin")
+CALENDAR_URI = QLIB_DIR.joinpath("calendars")
+INSTRUMENT_URI = QLIB_DIR.joinpath("instruments")
+FEATURE_URI = QLIB_DIR.joinpath("features")
 
 
 class TestStorage:
@@ -38,9 +39,10 @@ class TestStorage:
 
     def test_calendar_storage(self):
 
-        calendar = CalendarStorage(uri=CALENDAR_URI)
+        calendar = CalendarStorage(freq="day", future=False, uri=CALENDAR_URI)
         assert isinstance(calendar, Iterable), f"{calendar.__class__.__name__} is not Iterable"
         assert isinstance(calendar[:], Iterable), f"{calendar.__class__.__name__}.__getitem__(s: slice) is not Iterable"
+        assert isinstance(calendar.data, Iterable), f"{calendar.__class__.__name__}.data is not Iterable"
 
         print(f"calendar[1: 5]: {calendar[1:5]}")
         print(f"calendar[0]: {calendar[0]}")
@@ -80,11 +82,11 @@ class TestStorage:
 
         """
 
-        instrument = InstrumentStorage(uri=INSTRUMENT_URI)
+        instrument = InstrumentStorage(market="csi300", uri=INSTRUMENT_URI)
 
         assert isinstance(instrument, Iterable), f"{instrument.__class__.__name__} is not Iterable"
 
-        for inst, spans in instrument.items():
+        for inst, spans in instrument.data.items():
             assert isinstance(inst, str) and isinstance(
                 spans, Iterable
             ), f"{instrument.__class__.__name__} value is not Iterable"
@@ -149,11 +151,14 @@ class TestStorage:
 
         """
 
-        feature = FeatureStorage(uri=FEATURE_URI)
+        feature = FeatureStorage(instrument="SH600004", field="close", freq="day", uri=FEATURE_URI)
 
         assert isinstance(feature, Iterable), f"{feature.__class__.__name__} is not Iterable"
         with pytest.raises(IndexError):
             print(feature[0])
+        assert isinstance(
+            feature[815][1], (np.float, np.float32)
+        ), f"{feature.__class__.__name__}.__getitem__(i: int) error"
         assert len(feature[815:818]) == 3, f"{feature.__class__.__name__}.__getitem__(s: slice) error"
         print(f"feature[815: 818]: {feature[815: 818]}")
 
@@ -162,5 +167,5 @@ class TestStorage:
                 isinstance(_item, tuple) and len(_item) == 2
             ), f"{feature.__class__.__name__}.__iter__ item type error"
             assert isinstance(_item[0], int) and isinstance(
-                _item[1], (float, np.float, np.float32)
+                _item[1], (np.float, np.float32)
             ), f"{feature.__class__.__name__}.__iter__ value type error"