From 317357b50d2e54a2663e9fe5f14def3956423248 Mon Sep 17 00:00:00 2001 From: zhupr Date: Tue, 13 Apr 2021 10:47:01 +0800 Subject: [PATCH] Modify data.storage --- qlib/data/data.py | 92 +++++++++++++++-------------- qlib/data/storage/file_storage.py | 79 +++++++++++++++---------- qlib/data/storage/storage.py | 38 +++++++++--- tests/storage_tests/test_storage.py | 23 +++++--- 4 files changed, 141 insertions(+), 91 deletions(-) diff --git a/qlib/data/data.py b/qlib/data/data.py index 000bd1196..1a0ca616e 100644 --- a/qlib/data/data.py +++ b/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: diff --git a/qlib/data/storage/file_storage.py b/qlib/data/storage/file_storage.py index e2e5bd3e7..4090e3230 100644 --- a/qlib/data/storage/file_storage.py +++ b/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(" 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(" 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=" 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=" 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. """ diff --git a/tests/storage_tests/test_storage.py b/tests/storage_tests/test_storage.py index a70ce82ea..8ce3f5081 100644 --- a/tests/storage_tests/test_storage.py +++ b/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"