chinese-equity-quant/plugins/joinquant/wrapper_strategy.py

"""Generate a standalone JoinQuant wrapper strategy.

The module also defines default JoinQuant strategy hooks by executing the same
template with ``run1`` / ``target_shares`` defaults. That means this file can be
copied directly into JoinQuant for a quick smoke test, while the CLI can still
write a configured standalone file for a real run.
"""

from __future__ import annotations

import json
from pathlib import Path
from string import Template
from typing import Literal


WrapperMode = Literal["target_shares", "target_value"]


_WRAPPER_TEMPLATE = Template(
    r'''# Standalone JoinQuant target wrapper generated by chinese-equity-quant.
# Copy this file and the exported daily CSV target files into JoinQuant.
#
# The file loader is isolated in _read_target_file(). The default implementation
# uses JoinQuant's read_file API for uploaded files. If your JoinQuant runtime
# allows HTTP or another storage backend, replace only that function.

import csv
import io
import json


PORTFOLIO_NAME = "${portfolio_name}"
TARGET_MODE = "${mode}"
ALLOW_SHORT = ${allow_short}
TARGET_FILE_PREFIX = ""  # Optional uploaded-file prefix, for example "run1/"
_EMBEDDED_TARGETS = ${embedded_targets}


def initialize(context):
    set_benchmark("000300.XSHG")
    set_option("use_real_price", True)
    g.portfolio_name = PORTFOLIO_NAME
    g.target_mode = TARGET_MODE
    g.targets_by_date = {}
    run_daily(load_targets, time="before_open")
    run_daily(rebalance_at_open, time="open")
    run_daily(record_after_close, time="after_close")


def _today_text(context):
    return context.current_dt.strftime("%Y-%m-%d")


def _today_file_name(context):
    return context.current_dt.strftime("%Y%m%d") + ".csv"


def _read_target_file(file_name):
${embedded_target_read}
    data = read_file(TARGET_FILE_PREFIX + file_name)
    if isinstance(data, bytes):
        data = data.decode("utf-8")
    return data


def _load_target_rows(context):
    file_name = _today_file_name(context)
    text = _read_target_file(file_name)
    rows = list(csv.DictReader(io.StringIO(text)))
    clean_rows = []
    for row in rows:
        if row.get("portfolio_name") and row["portfolio_name"] != PORTFOLIO_NAME:
            continue
        jq_symbol = row.get("jq_symbol") or row.get("security") or row.get("symbol")
        if not jq_symbol:
            log.warn("Skipping target row with no jq_symbol: %s" % row)
            continue

        if TARGET_MODE == "target_shares":
            target = int(float(row.get("target_shares") or 0))
        elif TARGET_MODE == "target_value":
            target = float(row.get("target_value") or 0.0)
        else:
            raise ValueError("Unsupported TARGET_MODE: %s" % TARGET_MODE)

        if not ALLOW_SHORT and target < 0:
            log.warn(
                "SHORT_NOT_SUPPORTED clipping %s target from %s to 0" %
                (jq_symbol, target)
            )
            target = 0

        clean_rows.append({"jq_symbol": jq_symbol, "target": target, "raw": row})
    return clean_rows


def load_targets(context):
    date_text = _today_text(context)
    try:
        rows = _load_target_rows(context)
    except Exception as exc:
        log.error("Failed to load JoinQuant target file for %s: %s" % (date_text, exc))
        rows = []
    g.targets_by_date[date_text] = rows
    log.info("JOINQUANT_TARGET_LOAD|%s" % json.dumps({
        "date": date_text,
        "portfolio_name": PORTFOLIO_NAME,
        "target_mode": TARGET_MODE,
        "n_targets": len(rows),
    }, sort_keys=True))


def rebalance_at_open(context):
    date_text = _today_text(context)
    rows = g.targets_by_date.get(date_text, [])
    target_symbols = set()
    for row in rows:
        security = row["jq_symbol"]
        target_symbols.add(security)
        if TARGET_MODE == "target_shares":
            order_target(security, int(row["target"]))
        else:
            order_target_value(security, float(row["target"]))
        log.info("JOINQUANT_ORDER_SUBMIT|%s" % json.dumps({
            "date": date_text,
            "portfolio_name": PORTFOLIO_NAME,
            "jq_symbol": security,
            "target_mode": TARGET_MODE,
            "target": row["target"],
        }, sort_keys=True))

    for security in list(context.portfolio.positions.keys()):
        if security not in target_symbols:
            order_target(security, 0)
            log.info("JOINQUANT_ORDER_CLOSE|%s" % json.dumps({
                "date": date_text,
                "portfolio_name": PORTFOLIO_NAME,
                "jq_symbol": security,
            }, sort_keys=True))


def _position_records(context):
    records = []
    cash = float(context.portfolio.available_cash)
    total_value = float(context.portfolio.total_value)
    for security, position in context.portfolio.positions.items():
        records.append({
            "date": _today_text(context),
            "portfolio_name": PORTFOLIO_NAME,
            "jq_symbol": security,
            "position_shares": int(position.total_amount),
            "position_value": float(position.value),
            "cash": cash,
            "total_value": total_value,
        })
    return records


def _trade_records(context):
    records = []
    try:
        trades = get_trades()
    except Exception:
        trades = {}
    for trade_id, trade in trades.items():
        amount = int(getattr(trade, "amount", 0))
        price = float(getattr(trade, "price", 0.0))
        security = getattr(trade, "security", "")
        side = "buy" if amount >= 0 else "sell"
        records.append({
            "date": _today_text(context),
            "portfolio_name": PORTFOLIO_NAME,
            "jq_symbol": security,
            "order_id": str(getattr(trade, "order_id", trade_id)),
            "side": side,
            "filled_shares": amount,
            "fill_price": price,
            "trade_value": abs(amount * price),
            "trade_cost": float(getattr(trade, "commission", 0.0)),
            "raw_status": "filled",
        })
    return records


def record_after_close(context):
    date_text = _today_text(context)
    for record in _trade_records(context):
        log.info("JOINQUANT_FILL|%s" % json.dumps(record, sort_keys=True))
    for record in _position_records(context):
        log.info("JOINQUANT_POSITION|%s" % json.dumps(record, sort_keys=True))
    log.info("JOINQUANT_PNL|%s" % json.dumps({
        "date": date_text,
        "portfolio_name": PORTFOLIO_NAME,
        "cash": float(context.portfolio.available_cash),
        "total_value": float(context.portfolio.total_value),
    }, sort_keys=True))
'''
)


# Make this module itself usable as a JoinQuant strategy with defaults.
exec(_WRAPPER_TEMPLATE.substitute(
    portfolio_name="run1",
    mode="target_shares",
    allow_short="False",
    embedded_targets="{}",
    embedded_target_read="",
))


def render_wrapper_strategy(
    *,
    portfolio_name: str,
    mode: WrapperMode = "target_shares",
    allow_short: bool = False,
    embedded_targets: dict[str, str] | None = None,
) -> str:
    """Render the standalone JoinQuant wrapper strategy source."""
    if mode not in {"target_shares", "target_value"}:
        raise ValueError("mode must be 'target_shares' or 'target_value'")
    embedded_targets = embedded_targets or {}
    embedded_target_read = ""
    if embedded_targets:
        embedded_target_read = (
            "    if file_name in _EMBEDDED_TARGETS:\n"
            "        return _EMBEDDED_TARGETS[file_name]\n"
        )
    return _WRAPPER_TEMPLATE.substitute(
        portfolio_name=portfolio_name,
        mode=mode,
        allow_short="True" if allow_short else "False",
        embedded_targets=json.dumps(embedded_targets, ensure_ascii=False, indent=4),
        embedded_target_read=embedded_target_read,
    )


def _load_embedded_targets(targets_dir: str | Path | None) -> dict[str, str]:
    if targets_dir is None:
        return {}
    root = Path(targets_dir)
    targets = {
        path.name: path.read_text(encoding="utf-8")
        for path in sorted(root.glob("*.csv"))
    }
    if not targets:
        raise ValueError(f"No CSV target files found under {root}")
    return targets


def write_wrapper_strategy(
    *,
    portfolio_name: str,
    mode: WrapperMode = "target_shares",
    out_path: str | Path,
    allow_short: bool = False,
    embedded_targets_dir: str | Path | None = None,
) -> Path:
    """Write a configured standalone JoinQuant wrapper strategy."""
    path = Path(out_path)
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(
        render_wrapper_strategy(
            portfolio_name=portfolio_name,
            mode=mode,
            allow_short=allow_short,
            embedded_targets=_load_embedded_targets(embedded_targets_dir),
        ),
        encoding="utf-8",
    )
    return path