无限易套利合约网格交易 秒级行情。加入锁概念（Pro增强版）

import time


from threading import Lock

from pythongo.base import BaseParams, BaseState, Field
from pythongo.classdef import KLineData, OrderData, TickData, TradeData
from pythongo.core import KLineStyleType
from pythongo.ui import BaseStrategy
from pythongo.utils import KLineGeneratorArb
import talib
from pythongo.utils import KLineGenerator


class Params(BaseParams):
    """参数映射模型"""
    exchange: str = Field(default="", title="交易所代码")
    instrument_id: str = Field(default="", title="合约代码")
    initial_price: float = Field(default=0.0, title="初始价格")
    buy_grid_size: float = Field(default=6, title="买入网格大小")  # 新增买入网格大小
    sell_grid_size: float = Field(default=6, title="卖出网格大小")  # 新增卖出网格大小
    max_position: int = Field(default=10, title="最大持仓手数")
    order_volume: int = Field(default=1, title="下单手数", ge=1)
    pay_up: int | float = Field(default=1, title="超价")
    kline_style: KLineStyleType = Field(default="M1", title="K 线周期")


class State(BaseState):
    """状态映射模型"""
    current_price: float = Field(default=0.0, title="当前价格")
    buy_count: int = Field(default=0, title="买入次数")
    sell_count: int = Field(default=0, title="卖出次数")
    position: int = Field(default=0, title="当前持仓")
    is_order_completed: bool = Field(default=False, title="订单是否完成")  # 新增字段来跟踪订单是否完成


class GridTradingStrategyWithPositionLimit(BaseStrategy):
    """交易所标准套利合约的带持仓限制的网格交易策略"""

    def __init__(self):
        super().__init__()
        self.params_map = Params()
        """参数表"""
        self.kline_generator: KLineGenerator = None
        self.state_map = State()
        """状态表"""

        self.order_id = None
        """报单 ID"""
        self.order_lock = Lock()  # 创建一个锁
        self.order_in_progress = False  # 新增一个标志，用于表示是否有订单正在处理中

    @property
    def main_indicator_data(self) -> dict[str, float]:
        """主图指标"""
        return {
            "ShortMA": self.short_ma[-1],  # 短期均线
            "LongMA": self.long_ma[-1]  # 长期均线
        }

    def on_tick(self, tick: TickData) -> None:
        """收到行情 tick 推送"""
        super().on_tick(tick)

        self.kline_generator.tick_to_kline(tick)

    def on_order_cancel(self, order: OrderData) -> None:
        """撤单推送回调"""
        super().on_order_cancel(order)
        self.order_id = None
        self.state_map.is_order_completed = False  # 撤单时标记订单未完成

    def on_trade(self, trade: TradeData, log: bool = False) -> None:
        """成交回调"""
        super().on_trade(trade, log)
        if self.order_id and self.order_id == trade.order_id:  # 判断成交的订单是否是之前发送的订单
            self.state_map.is_order_completed = True  # 成交时标记订单完成
            self.order_lock.release()  # 解锁
        self.order_id = None

        # # 更新持仓
        if trade.direction == "多":
            self.state_map.position += trade.volume
            self.params_map.initial_price = trade.price  # 买入成功后更新初始价格
            self.state_map.buy_count += 1
        elif trade.direction == "空":
            self.state_map.position -= trade.volume
            self.params_map.initial_price = trade.price  # 卖出成功后更新初始价格
            self.state_map.sell_count += 1

    def on_start(self):
        self.kline_generator = KLineGeneratorArb(
            callback=self.callback,
            exchange=self.params_map.exchange,
            instrument_id=self.params_map.instrument_id,
            style=self.params_map.kline_style,
            real_time_callback=self.real_time_callback
        )
        self.kline_generator.push_history_data()

        super().on_start()

        for instrument_id in self.kline_generator.instruments:
            """订阅单腿合约"""
            self.sub_market_data(
                exchange=self.params_map.exchange,
                instrument_id=instrument_id
            )

    def on_stop(self):
        super().on_stop()

        for instrument_id in self.kline_generator.instruments:
            """取消订阅单腿合约"""
            self.unsub_market_data(
                exchange=self.params_map.exchange,
                instrument_id=instrument_id
            )

    def callback(self, kline: KLineData):
        """接受 K 线回调"""

        self.calc_double_ma()
        self.check_grid_triggers()
        self.widget.recv_kline({
            "kline": kline,
            **self.main_indicator_data,

        })

        if self.trading:
            self.update_status_bar()

    def real_time_callback(self, kline: KLineData) -> None:

        self.calc_double_ma()
        """使用收到的实时推送 K 线来更新价格并检查触发条件"""
        # 获取真实持仓
        position_info = self.get_position(self.params_map.instrument_id)  # 假设存在这样一个获取持仓的方法
        # 更新状态中的持仓信息
        self.state_map.position = position_info.net_position  # 根据实际的持仓数据结构进行赋值
        self.state_map.current_price = kline.close
        self.check_grid_triggers()
        self.widget.recv_kline({
            "kline": kline,
            **self.main_indicator_data,

        })

    def check_grid_triggers(self):
        """检查网格交易的触发条件"""
        price_change = self.state_map.current_price - self.params_map.initial_price

        # 当价格上涨超过卖出网格大小且持仓未达到最大时卖出
        if price_change >= self.params_map.sell_grid_size and abs(
                self.state_map.position) < self.params_map.max_position:
            self.sell()

        # 当价格下跌超过买入网格大小且持仓未达到最大负数时买入
        elif price_change <= -self.params_map.buy_grid_size and self.state_map.position < self.params_map.max_position:
            self.buy()

    def sell(self):
        """执行卖出操作"""
        try:
            if not self.order_lock.acquire(timeout=5):  # 设置超时时间为 5 秒
                self.output("获取锁超时，放弃本次卖出操作")
                return

            if self.state_map.position < self.params_map.max_position and not self.order_in_progress:
                self.order_in_progress = True  # 设置订单正在处理中
                self.order_id = self.send_order(
                    exchange=self.params_map.exchange,
                    instrument_id=self.params_map.instrument_id,
                    volume=self.params_map.order_volume,
                    price=(price := self.state_map.current_price - self.params_map.pay_up),
                    order_direction="sell"
                )

                self.output(f"卖出信号价格: {price}")
                while not self.state_map.is_order_completed:  # 等待订单完成
                    time.sleep(1)  # 可以根据需要调整等待间隔
        except Exception:
            self.output("获取锁超时，放弃本次卖出操作")
        finally:
            self.order_lock.release()  # 确保释放锁
            if self.order_in_progress:  # 如果之前的订单处理完成，重置标志
                self.order_in_progress = False

    def buy(self):
        """执行买入操作"""
        try:
            if not self.order_lock.acquire(timeout=5):  # 设置超时时间为 5 秒
                self.output("获取锁超时，放弃本次买入操作")
                return

            if self.state_map.position > -self.params_map.max_position and not self.order_in_progress:
                self.order_in_progress = True  # 设置订单正在处理中
                self.order_id = self.send_order(
                    exchange=self.params_map.exchange,
                    instrument_id=self.params_map.instrument_id,
                    volume=self.params_map.order_volume,
                    price=(price := self.state_map.current_price + self.params_map.pay_up),
                    order_direction="buy"
                )

                self.output(f"买入信号价格: {price}")
                while not self.state_map.is_order_completed:  # 等待订单完成
                    time.sleep(1)  # 可以根据需要调整等待间隔
        except Exception:
            self.output("获取锁超时，放弃本次买入操作")
        finally:
            self.order_lock.release()  # 确保释放锁
            if self.order_in_progress:  # 如果之前的订单处理完成，重置标志
                self.order_in_progress = False

    def calc_bollinger_bands(self) -> None:
        """计算布林带指标"""
        close = self.kline_generator.producer.close
        upper, middle, lower = talib.BBANDS(close, timeperiod=20, nbdevup=2, nbdevdn=2)
        self.upper_band = upper
        self.middle_band = middle
        self.lower_band = lower

    def calc_double_ma(self) -> None:
        """计算双均线"""
        close = self.kline_generator.producer.close
        # 假设短期均线为 5 周期，长期均线为 20 周期
        self.short_ma = talib.SMA(close, timeperiod=5)
        self.long_ma = talib.SMA(close, timeperiod=20)