py, update sina data api

yssource · Feb 7, 2022 · 19b937e · 19b937e
1 parent fdaf76f
commit 19b937e
Show file tree

Hide file tree

Showing 4 changed files with 49 additions and 8 deletions.
diff --git a/abquant/apis/base.py b/abquant/apis/base.py
@@ -245,5 +245,7 @@ def get_security_info(
 
 def get_realtime_quotes(codes: Iterable[str]) -> pd.DataFrame:
     import abquant.data.ts_api as ts
+    from abquant.utils.ts import order_book_id_2_tushare_code
 
+    codes = [order_book_id_2_tushare_code(code) for code in codes]
     return ts.get_realtime_quotes(codes)
diff --git a/abquant/data/ts_api.py b/abquant/data/ts_api.py
@@ -9,8 +9,10 @@
     US_LIVE_DATA_COLS,
     LIVE_DATA_COLS,
 )
+from abquant.utils.logger import user_log as ulog
 import re
-from urllib.request import urlopen, Request
+import requests
+from requests.exceptions import Timeout
 
 
 def get_realtime_quotes(symbols=None) -> pd.DataFrame:
@@ -61,11 +63,17 @@ def get_realtime_quotes(symbols=None) -> pd.DataFrame:
     else:
         symbols_list = code_to_symbol(symbols)
     symbols_list = symbols_list[:-1] if len(symbols_list) > 8 else symbols_list
-    request = Request(
-        LIVE_DATA_URL % (P_TYPE["http"], DOMAINS["sinahq"], random(), symbols_list)
-    )
-    text = urlopen(request, timeout=10).read()
-    text = text.decode("GBK")
+    url = LIVE_DATA_URL % (P_TYPE["http"], DOMAINS["sinahq"], random(), symbols_list)
+    headers = {"host": "hq.sinajs.cn", "referer": "https://finance.sina.com.cn/"}
+
+    text = ""
+    try:
+        resp = requests.get(url, headers=headers, timeout=10)
+        text = resp.text
+    except Timeout:
+        # ulog.warning(f"URL: {url} \n Timeout has been raised.")
+        return pd.DataFrame()
+
     reg = re.compile(r'\="(.*?)\";')
     data = reg.findall(text)
     regSym = re.compile(r"(?:sh|sz|gb_)(.*?)\=")
@@ -74,7 +82,7 @@ def get_realtime_quotes(symbols=None) -> pd.DataFrame:
     syms_list = []
     for index, row in enumerate(data):
         if len(row) > 1:
-            data_list.append([astr for astr in row.split(",")][:len(LIVE_DATA_COLS)])
+            data_list.append([astr for astr in row.split(",")][: len(LIVE_DATA_COLS)])
             syms_list.append(syms[index])
     if len(syms_list) == 0:
         return pd.DataFrame()

diff --git a/abquant/utils/ts.py b/abquant/utils/ts.py
@@ -129,3 +129,7 @@ def random(n: int = 13) -> str:
     start = 10 ** (n - 1)
     end = (10 ** n) - 1
     return str(randint(start, end))
+
+
+def order_book_id_2_tushare_code(order_book_id: str) -> str:
+    return order_book_id.split(".")[0]
diff --git a/tests/apis/test_base.py b/tests/apis/test_base.py
@@ -1,5 +1,10 @@
 from pyabquant import PyAbquant, FQ_TYPE, INSTRUMENT_TYPE  # type: ignore
-from abquant.apis.base import get_price, get_all_securities, get_security_info
+from abquant.apis.base import (
+    get_price,
+    get_all_securities,
+    get_security_info,
+    get_realtime_quotes,
+)
 import pandas as pd
 from abquant.utils.logger import user_log as ulog
 
@@ -143,3 +148,25 @@ def test_get_security_info_index_type():
     ulog.debug(actual)
     assert isinstance(actual, pd.DataFrame)
     assert actual.display_name == "创业板指"
+
+
+def test_get_realtime_quotes():
+    codes = ["000001.XSHE", "600000.XSHG", "300001.XSHE"]
+    names = ["平安银行", "浦发银行", "特锐德"]
+    actual = get_realtime_quotes(codes)
+    ulog.debug(actual)
+    ulog.debug(actual["name"])
+    ulog.debug(pd.Series(names))
+    assert isinstance(actual, pd.DataFrame)
+    assert actual.name.equals(pd.Series(names))
+
+
+def test_get_realtime_quotes_index():
+    codes = ["sh000001", "sz399001", "sh000300", "sh000016", "sz399005", "sz399006"]
+    names = ["上证指数", "深证成指", "沪深300", "上证50", "中小100", "创业板指"]
+    actual = get_realtime_quotes(codes)
+    ulog.debug(actual)
+    ulog.debug(actual["name"])
+    ulog.debug(pd.Series(names))
+    assert isinstance(actual, pd.DataFrame)
+    assert actual.name.equals(pd.Series(names))