fix: possible bug causing missing calendar_list data

scripts/data_collector/utils.py

@@ -3,10 +3,12 @@
 
 import re
 import copy
+import datetime
 import importlib
 import time
 import bisect
 import pickle
+import random
 import requests
 import functools
 from pathlib import Path
@@ -23,7 +25,7 @@ from bs4 import BeautifulSoup
 
 HS_SYMBOLS_URL = "http://app.finance.ifeng.com/hq/list.php?type=stock_a&class={s_type}"
 
-CALENDAR_URL_BASE = "http://push2his.eastmoney.com/api/qt/stock/kline/get?secid={market}.{bench_code}&fields1=f1%2Cf2%2Cf3%2Cf4%2Cf5&fields2=f51%2Cf52%2Cf53%2Cf54%2Cf55%2Cf56%2Cf57%2Cf58&klt=101&fqt=0&beg=19900101&end=20991231"
+CALENDAR_URL_BASE = "http://push2his.eastmoney.com/api/qt/stock/kline/get?secid={market}.{bench_code}&fields1=f1%2Cf2%2Cf3%2Cf4%2Cf5&fields2=f51%2Cf52%2Cf53%2Cf54%2Cf55%2Cf56%2Cf57%2Cf58&klt=101&fqt=0"
 SZSE_CALENDAR_URL = "http://www.szse.cn/api/report/exchange/onepersistenthour/monthList?month={month}&random={random}"
 
 CALENDAR_BENCH_URL_MAP = {
@@ -38,6 +40,24 @@ CALENDAR_BENCH_URL_MAP = {
     "BR_ALL": "^BVSP",
 }
 
+CHROME_UA_POOL = [
+    # Windows
+    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
+    "AppleWebKit/537.36 (KHTML, like Gecko) "
+    "Chrome/120.0.0.0 Safari/537.36",
+    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
+    "AppleWebKit/537.36 (KHTML, like Gecko) "
+    "Chrome/121.0.6167.85 Safari/537.36",
+    # macOS
+    "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) "
+    "AppleWebKit/537.36 (KHTML, like Gecko) "
+    "Chrome/121.0.0.0 Safari/537.36",
+    # Linux
+    "Mozilla/5.0 (X11; Linux x86_64) "  # pylint: disable=W1404
+    "AppleWebKit/537.36 (KHTML, like Gecko) "
+    "Chrome/120.0.0.0 Safari/537.36",
+]
+
 _BENCH_CALENDAR_LIST = None
 _ALL_CALENDAR_LIST = None
 _HS_SYMBOLS = None
@@ -51,6 +71,16 @@ _CALENDAR_MAP = {}
 MINIMUM_SYMBOLS_NUM = 3900
 
 
+def build_headers():
+    return {
+        "User-Agent": random.choice(CHROME_UA_POOL),
+        "Accept": "application/json,text/plain,*/*",
+        "Accept-Language": "zh-CN,zh;q=0.9",
+        "Referer": "https://quote.eastmoney.com/",
+        "Connection": "keep-alive",
+    }
+
+
 def get_calendar_list(bench_code="CSI300") -> List[pd.Timestamp]:
     """get SH/SZ history calendar list
 
@@ -66,9 +96,38 @@ def get_calendar_list(bench_code="CSI300") -> List[pd.Timestamp]:
 
     logger.info(f"get calendar list: {bench_code}......")
 
-    def _get_calendar(url):
-        _value_list = requests.get(url, timeout=None).json()["data"]["klines"]
-        return sorted(map(lambda x: pd.Timestamp(x.split(",")[0]), _value_list))
+    def _get_calendar(url, max_retry=3):
+        session = requests.Session()
+        session.headers.update(build_headers())
+        current_datetime = datetime.datetime.now()
+        cur_year = current_datetime.year
+        res_list = []
+        failed_years = []
+        for year in range(2000, cur_year + 1):
+            start = f"{year}0101"
+            end = f"{year}1231"
+            formatted_url = url + f"&beg={start}&end={end}".format(start=start, end=end)
+            for attempt in range(max_retry):
+                try:
+                    resp = session.get(formatted_url, timeout=10)
+                    resp.raise_for_status()
+
+                    data = resp.json().get("data")
+                    if not data or "klines" not in data:
+                        raise ValueError("missing klines")
+
+                    res_list.extend(pd.Timestamp(x.split(",")[0]) for x in data["klines"])
+                    break
+
+                except Exception as e:
+                    time.sleep(random.uniform(0.8, 1.5))
+            else:
+                failed_years.append(year)
+
+        if failed_years:
+            logger.warning(f"Calendar incomplete, failed years: {failed_years}")
+
+        return sorted(set(res_list))
 
     calendar = _CALENDAR_MAP.get(bench_code, None)
     if calendar is None:
@@ -246,7 +305,11 @@ def get_hs_stock_symbols() -> list:
 
         # Add suffix after the stock code to conform to yahooquery standard, otherwise the data will not be fetched.
         _symbols = [
-            _symbol + ".ss" if _symbol.startswith("6") else _symbol + ".sz" if _symbol.startswith(("0", "3")) else None
+            (
+                _symbol + ".ss"
+                if _symbol.startswith("6")
+                else _symbol + ".sz" if _symbol.startswith(("0", "3")) else None
+            )
             for _symbol in _symbols
         ]
         _symbols = [_symbol for _symbol in _symbols if _symbol is not None]
@@ -357,7 +420,14 @@ def get_us_stock_symbols(qlib_data_path: [str, Path] = None) -> list:
             s_ = s_.strip("*")
             return s_
 
-        _US_SYMBOLS = sorted(set(map(_format, filter(lambda x: len(x) < 8 and not x.endswith("WS"), _all_symbols))))
+        _US_SYMBOLS = sorted(
+            set(
+                map(
+                    _format,
+                    filter(lambda x: len(x) < 8 and not x.endswith("WS"), _all_symbols),
+                )
+            )
+        )
 
     return _US_SYMBOLS
 
@@ -471,7 +541,10 @@ def get_en_fund_symbols(qlib_data_path: [str, Path] = None) -> list:
             raise ValueError("request error")
         try:
             _symbols = []
-            for sub_data in re.findall(r"[\[](.*?)[\]]", resp.content.decode().split("= [")[-1].replace("];", "")):
+            for sub_data in re.findall(
+                r"[\[](.*?)[\]]",
+                resp.content.decode().split("= [")[-1].replace("];", ""),
+            ):
                 data = sub_data.replace('"', "").replace("'", "")
                 # TODO: do we need other information, like fund_name from ['000001', 'HXCZHH', '华夏成长混合', '混合型', 'HUAXIACHENGZHANGHUNHE']
                 _symbols.append(data.split(",")[0])
@@ -652,7 +725,11 @@ def get_instruments(
     """
     _cur_module = importlib.import_module("data_collector.{}.collector".format(market_index))
     obj = getattr(_cur_module, f"{index_name.upper()}Index")(
-        qlib_dir=qlib_dir, index_name=index_name, freq=freq, request_retry=request_retry, retry_sleep=retry_sleep
+        qlib_dir=qlib_dir,
+        index_name=index_name,
+        freq=freq,
+        request_retry=request_retry,
+        retry_sleep=retry_sleep,
     )
     getattr(obj, method)()
 
@@ -660,7 +737,10 @@ def get_instruments(
 def _get_all_1d_data(_date_field_name: str, _symbol_field_name: str, _1d_data_all: pd.DataFrame):
     df = copy.deepcopy(_1d_data_all)
     df.reset_index(inplace=True)
-    df.rename(columns={"datetime": _date_field_name, "instrument": _symbol_field_name}, inplace=True)
+    df.rename(
+        columns={"datetime": _date_field_name, "instrument": _symbol_field_name},
+        inplace=True,
+    )
     df.columns = list(map(lambda x: x[1:] if x.startswith("$") else x, df.columns))
     return df