main
  1#!/usr/bin/env python
  2# -*- coding: utf-8 -*-
  3import re
  4
  5from pyrogram.client import Client
  6from pyrogram.types import Message, ReplyParameters
  7
  8from config import PREFIX, TZ, cache
  9from messages.parser import parse_msg
 10from messages.progress import modify_progress
 11from messages.sender import send2tg
 12from messages.utils import equal_prefix, startswith_prefix
 13from price.binance import binance_supported, get_binance_price
 14from price.coinmarketcap import cmc_convert_price, cmc_supported, get_cmc_price
 15from price.okx import get_okx_price, okx_supported
 16from price.tradingview import get_tradingview_price, tradingview_supported
 17
 18HELP = f"""
 19💵**查询价格**
 20使用说明:
 21- `{PREFIX.PRICE}` + Symbol + [@Interval]
 22- 或`{PREFIX.CRYPTO}` 仅查询加密货币市场
 23- 或`{PREFIX.STOCK}` 仅查询股票市场
 24
 25其中symbol(大小写不限)支持如下类别:
 261. 加密货币, 如 `BTC`
 272. 股票, 如 `AAPL` / `SPX` (A股, 港股, 美股)
 283. 汇率, 如 `USD CNY` (中间有空格)
 29
 30K线Interval (可选):
 31- 加密货币(默认30m)
 321m,3m,5m,15m,30m,1h,2h,4h,6h,8h,12h,1D,3D,1W,1M
 33- 股票(默认5m)
 341m,3m,5m,15m,30m,45m,1h,2h,3h,4h,1D,1W,1M,3M,6M,1Y
 35
 36说明:
 37- 加密货币支持币种代码(BTC), 币种名称(bitcoin), 或交易对(BTCUSDC).
 38- 此外在Symbol前添加数字可以计算对应数量的价值。当前仅支持对加密货币和法币汇率进行计算。
 39- 有些代码同时被加密货币和股票市场使用而产生冲突。例如`SPX`是标普500指数, 同时也是一个山寨币。此时可以使用`{PREFIX.CRYPTO}`或`{PREFIX.STOCK}`来限定市场。
 40
 41示例:
 421. 查询加密货币价格
 43- 对于Binance和OKX支持的币种, 还会返回K线图(默认30m)
 44- `{PREFIX.PRICE} BTC`
 45- `{PREFIX.PRICE} ethereum`
 46- `{PREFIX.PRICE} DOGEUSDT`
 47- `{PREFIX.PRICE} BTC @4h`
 48
 492. 查询股票价格:
 50- 默认返回Interval为5m的K线图
 51- `{PREFIX.PRICE}` AAPL 或 NASDAQ:AAPL
 52- `{PREFIX.PRICE}` SPX 或 SP:SPX
 53- `{PREFIX.PRICE}` 000001 或 SSE:000001
 54- `{PREFIX.PRICE} AAPL @1m`
 55
 563. 查询汇率:
 57- `{PREFIX.PRICE} USD CNY`
 58- `{PREFIX.PRICE} BTC CNY`
 59- `{PREFIX.PRICE} DOGE BTC`
 60
 614. 计算价值:
 62- `{PREFIX.PRICE} 1.5 BTC` (默认计算美元价值)
 63- `{PREFIX.PRICE} 1.5 BTC CNY`
 64- `{PREFIX.PRICE} 3000 JPY CNY`
 65"""
 66
 67
 68async def get_asset_price(client: Client, message: Message, **kwargs):
 69    """Get asset price."""
 70    info = parse_msg(message)
 71    # send docs if message == "/price"
 72    if equal_prefix(info["text"], prefix=[PREFIX.PRICE, PREFIX.CRYPTO, PREFIX.STOCK]):
 73        await send2tg(client, message, texts=HELP, **kwargs)
 74        return
 75
 76    if not startswith_prefix(info["text"], prefix=[PREFIX.PRICE, PREFIX.CRYPTO, PREFIX.STOCK]):
 77        return
 78    prefix = info["text"].split(" ")[0]
 79    crypto_only = prefix == PREFIX.CRYPTO
 80    stock_only = prefix == PREFIX.STOCK
 81    text = info["text"].removeprefix(prefix).strip()
 82    # these patterns should use CoinMarketCap API
 83    # some coin has "$" in symbol, so we need to match it
 84    pattern_1 = r"^([\d.]+)\s+([$\dA-Za-z]+)\s+([$\dA-Za-z]+)$"  # match "1.5 BTC CNY"
 85    pattern_2 = r"^([\d.]+)\s+([$\dA-Za-z]+)$"  # match "1.5 BTC"
 86    pattern_3 = r"^([$\dA-Za-z]+)\s+([$\dA-Za-z]+)$"  # match "BTC CNY"
 87    amount, base, quote = 0, "", ""
 88    if matched := re.search(pattern_1, text, re.IGNORECASE):
 89        amount = float(matched.group(1))
 90        base = matched.group(2)
 91        quote = matched.group(3)
 92    elif matched := re.search(pattern_2, text, re.IGNORECASE):
 93        amount = float(matched.group(1))
 94        base = matched.group(2)
 95        quote = "USD"
 96    elif matched := re.search(pattern_3, text, re.IGNORECASE):
 97        amount = 1
 98        base = matched.group(1)
 99        quote = matched.group(2)
100    if not stock_only and amount > 0 and base and quote and (msg := await cmc_convert_price(amount, base, quote)):
101        await send2tg(client, message, texts=msg, **kwargs)
102        return
103    # match interval: "BTC @1m" or "000001 @15m"
104    if matched := re.search(r"^([$\dA-Za-z]+)\s+@(\d+[A-Za-z])$", text, re.IGNORECASE):
105        symbol = matched.group(1)
106        interval = matched.group(2)
107    else:  # match single symbol: "BTC" / "AAPL" / "SPX" / "000001"
108        symbol = text
109        interval = None
110    categories = await match_symbol_category(symbol, crypto_only=crypto_only, stock_only=stock_only)
111    if not categories:
112        await send2tg(client, message, texts=f"不支持此Symbol: {symbol.upper()}\n{HELP}", **kwargs)
113        return
114    msg = f"🔍查询价格: {symbol.upper()}"
115    if kwargs.get("show_progress"):
116        res = await send2tg(client, message, texts=msg, **kwargs)
117        kwargs["progress"] = res[0]
118    if warnings := categories.get("warnings"):
119        await modify_progress(text=warnings, **kwargs)
120    # Tradingview
121    if not crypto_only and categories.get("tradingview") and (data := await get_tradingview_price(symbol, interval, **kwargs)):
122        await client.send_photo(
123            chat_id=info["cid"],
124            photo=data["url"],
125            caption=f"[{data['symbol']}](https://www.tradingview.com/chart/?symbol={data['symbol']}) @{data['interval']} ({TZ})",
126            reply_parameters=ReplyParameters(message_id=info["mid"]),
127        )
128        await modify_progress(del_status=True, **kwargs)
129        return
130
131    # Belows are only for crypto market
132    if stock_only:
133        await modify_progress(del_status=True, **kwargs)
134        return
135    # Binance & OKX will return klines chart
136    if (res := await get_binance_price(symbol, interval)) or (res := await get_okx_price(symbol, interval)):
137        await send2tg(client, message, **res, **kwargs)
138        await modify_progress(del_status=True, **kwargs)
139        return
140    # other crypto assets supported by CoinMarketCap
141    if res := await get_cmc_price(text):
142        await send2tg(client, message, texts=res, **kwargs)
143    await modify_progress(del_status=True, **kwargs)
144
145
146@cache.memoize(ttl=3600)
147async def match_symbol_category(symbol: str = "", *, crypto_only: bool = False, stock_only: bool = False) -> dict[str, str]:
148    category = {}
149    # Crypto market
150    if not stock_only:
151        if cmc := await cmc_supported(symbol):  # {"symbol": "BTC"} or {"slug": "bitcoin"}
152            category["cmc"] = cmc.get("symbol", cmc.get("slug"))
153            category["crypto"] = cmc.get("symbol", cmc.get("slug"))
154        # okx
155        okx_symbol, _ = await okx_supported(symbol)  # (symbol, instId)
156        if okx_symbol:
157            category["okx"] = okx_symbol
158            category["crypto"] = okx_symbol
159        binance_symbol, _ = await binance_supported(symbol)  # (symbol, market)
160        if binance_symbol:
161            category["binance"] = binance_symbol
162            category["crypto"] = binance_symbol
163
164    # Stock market
165    tv_symbols = []
166    if not crypto_only and (tradingview := await tradingview_supported(symbol)):  # [("SSE:000001", "china"), ("SZSE:000001", "china")]
167        tv_symbols = [x[0] for x in tradingview]  # ["SSE:000001", "SZSE:000001"]
168        category["tradingview"] = ", ".join(tv_symbols)
169    # skip some crypto ETF (e.g. Grayscale Bitcoin Mini Trust use symbol "AMEX:BTC")
170    if category.get("crypto") and category.get("tradingview"):
171        exchange, coin = tradingview[0][0].split(":")  # type: ignore
172        if exchange == "AMEX" and category.get("crypto", "").startswith(coin):  # hit crypto ETF
173            tradingview = []
174            del category["tradingview"]
175
176    # if tradingview has multiles symbols
177    if len(tv_symbols) > 1:
178        msg = f"⚠️**{symbol.upper()}**代码重复:\n"
179        msg += f"股票: **{category['tradingview']}**\n"
180        msg += f"本次查询: **{tv_symbols[0]}**\n"
181        msg += f"查询其他请使用完整代码:\n`{PREFIX.PRICE} {'/ '.join(tv_symbols[1:])}`"
182        category["warnings"] = msg
183
184    # if symbol matches tradingview & crypto categories
185    if category.get("crypto") and category.get("tradingview"):
186        msg = f"⚠️**{symbol.upper()}**代码重复:\n"
187        msg += f"股票: **{category['tradingview']}**\n"
188        msg += f"加密货币: **{category['crypto']}**\n"
189        msg += f"默认查询股票: **{tv_symbols[0]}**\n"
190        msg += f"使用 `{PREFIX.CRYPTO}` 或 `{PREFIX.STOCK}` 限定市场"
191        category["warnings"] = msg
192
193    return category