add last_price

Author: riiy

src/emtl/core.py

@@ -1,3 +1,4 @@
+import os
 import re
 from random import SystemRandom
 from typing import Any
@@ -12,6 +13,7 @@
 from .const import _base_headers
 from .const import _urls
 from .utils import emt_trade_encrypt
+from .utils import get_float
 from .utils import get_logger
 
 logger = get_logger(__name__)
@@ -20,6 +22,23 @@
 _em_validate_key = ""
 
 
+def _query_snapshot(symbol_code: str, market: str) -> Optional[dict]:
+    url = "https://emhsmarketwg.eastmoneysec.com/api/SHSZQuoteSnapshot"
+    params = {"id": symbol_code.strip(), "market": market}
+    headers = _base_headers.copy()
+    resp = session.get(url, params=params, headers=headers)
+    _check_resp(resp)
+    return resp.json()
+
+
+def get_last_price(symbol_code: str, market: str) -> float:
+    ret = _query_snapshot(symbol_code, market)
+    if ret is None or "status" not in ret or ret["status"] != 0:
+        return float("nan")
+
+    return get_float(ret["realtimequote"], "currentPrice")
+
+
 def _check_resp(resp: Response):
     if resp.status_code != 200:
         logger.error(f"request {resp.url} fail, code={resp.status_code}, response={resp.text}")
@@ -33,9 +52,13 @@ def _query_something(tag: str, req_data: Optional[dict] = None) -> Optional[Dict
     :param req_data: 请求提交数据,可选
     :return:
     """
-    assert _em_validate_key, "em_validatekey is empty"
+    if not _em_validate_key:
+        validate_key = login()
+    else:
+        validate_key = _em_validate_key
+    logger.info(validate_key)
     assert tag in _urls, f"{tag} not in url list"
-    url = _urls[tag] + _em_validate_key
+    url = _urls[tag] + validate_key
     if req_data is None:
         req_data = {
             "qqhs": 100,
@@ -46,6 +69,7 @@ def _query_something(tag: str, req_data: Optional[dict] = None) -> Optional[Dict
     logger.debug(f"(tag={tag}), (data={req_data}), (url={url})")
     resp = session.post(url, headers=headers, data=req_data)
     _check_resp(resp)
+    logger.info(resp.text)
     return resp.json()
 
 
@@ -79,7 +103,7 @@ def _get_em_validate_key():
         return _em_validatekey
 
 
-def login(username: str, password: str, duration: int = 30) -> Optional[str]:
+def login(username: str = "", password: str = "", duration: int = 30) -> Optional[str]:
     """登录接口.
 
     :param str username: 用户名
@@ -88,6 +112,10 @@ def login(username: str, password: str, duration: int = 30) -> Optional[str]:
     :type duration: int, optional
     :return:
     """
+    if not username:
+        username = os.getenv("EM_USERNAME", "")
+    if not password:
+        password = os.getenv("EM_PASSWORD", "")
     random_num, code = _get_captcha_code()
     headers = _base_headers.copy()
     headers["X-Requested-With"] = "XMLHttpRequest"
@@ -199,5 +227,6 @@ def insert_order(stock_code, trade_type, market: str, price: float, amount: int)
 def cancel_order(code: str):
     data = {"revokes": code.strip()}
     resp = _query_something("cancel_order", req_data=data)
+    logger.info(resp)
     if resp:
         return resp

src/emtl/tests/test_core.py

@@ -1,7 +1,9 @@
 import datetime
+import math
 import os
 
 from emtl.core import cancel_order
+from emtl.core import get_last_price
 from emtl.core import insert_order
 from emtl.core import query_funds_flow
 from emtl.core import query_history_orders
@@ -80,3 +82,7 @@ def test_cancel_order():
     resp = cancel_order("000002")
     assert resp
     assert resp["Status"] in (0, -1)
+
+
+def test_last_price():
+    assert math.isnan(get_last_price("000001", "SH"))

src/emtl/utils.py

@@ -1,5 +1,6 @@
 import base64
 import logging
+import math
 
 from cryptography.hazmat.primitives import serialization
 from cryptography.hazmat.primitives.asymmetric import padding
@@ -27,3 +28,19 @@ def emt_trade_encrypt(content: str) -> str:
     _pub_key: rsa.RSAPublicKey = serialization.load_pem_public_key(rsa_public_key.encode("utf-8"))  # type:ignore
     encrypt_text = _pub_key.encrypt(content.encode(), padding.PKCS1v15())
     return base64.b64encode(encrypt_text).decode("utf-8")
+
+
+def double_equal(a, b) -> bool:
+    return math.fabs(a - b) < 1e-6
+
+
+def get_float(data: dict, key: str) -> float:
+    if v := data[key].strip():
+        return float(v)
+    return 0.0
+
+
+def get_int(data: dict, key: str) -> int:
+    if v := data[key].strip():
+        return int(v)
+    return 0