fix bug on TRA dataset (microsoft#1135)

* fix bug on TRA dataset solve issue "qrun TRA model error (microsoft#1062)" * apply black pylint
AIGeniusInstitute · Jul 7, 2022 · e62684e · e62684e
1 parent 8a5efda
commit e62684e
Showing 1 changed file with 8 additions and 2 deletions.
diff --git a/qlib/contrib/data/dataset.py b/qlib/contrib/data/dataset.py
@@ -203,8 +203,14 @@ def setup_data(self, handler_kwargs: dict = None, **kwargs):
 
     def _prepare_seg(self, slc, **kwargs):
         fn = _get_date_parse_fn(self._index[0][1])
-        start_date = fn(slc.start)
-        end_date = fn(slc.stop)
+        if isinstance(slc, slice):
+            start, stop = slc.start, slc.stop
+        elif isinstance(slc, (list, tuple)):
+            start, stop = slc
+        else:
+            raise NotImplementedError(f"This type of input is not supported")
+        start_date = pd.Timestamp(fn(start))
+        end_date = pd.Timestamp(fn(stop))
         obj = copy.copy(self)  # shallow copy
         # NOTE: Seriable will disable copy `self._data` so we manually assign them here
         obj._data = self._data  # reference (no copy)