文件预览

init.py

查看 paper-test2 技能包中的文件内容。

文件内容

scripts/init.py

"""First-run setup. Idempotent. Triggered by lifecycle.start() when .venv/ missing.

Steps performed by `run()`:
  1. mkdir -p config data state logs
  2. materialize templated configs into config/ (skipping any pre-existing file)
  3. generate state/token (32-byte hex, mode 0600) if missing
  4. python3 -m venv .venv (if venv python missing)
  5. venv pip install -e .

`materialize_configs` and `ensure_token` are safe to call independently;
`ensure_venv` requires pip + an internet connection on first call.
"""
from __future__ import annotations
import os
import secrets
import subprocess
import sys
from pathlib import Path

from ._venv_paths import venv_has_python, venv_pip, venv_python


SKILL_ROOT_DEFAULT = Path(__file__).resolve().parent.parent
TEMPLATES_DIR_REL = "templates/config"
CONFIG_FILES = ("main.yaml", "accounts.yaml", "markets.yaml", "symbols.yaml")
DIRS = ("config", "data", "state", "logs")


def _ensure_dirs(skill_root: Path) -> None:
    for sub in DIRS:
        (skill_root / sub).mkdir(parents=True, exist_ok=True)


def _templates_dir() -> Path:
    """Templates ship with the skill, not in the user's workspace.

    Resolved relative to this file so the function is cwd-independent.
    """
    return Path(__file__).resolve().parent.parent / TEMPLATES_DIR_REL


def is_first_run(skill_root: Path = SKILL_ROOT_DEFAULT) -> bool:
    """Return True if config/accounts.yaml has not been written yet.

    Use this to detect whether initial-cash confirmation is needed before
    the first service start. Once the service has run, accounts are seeded
    into SQLite — changing initial_cash in accounts.yaml has no effect.
    """
    return not (skill_root / "config" / "accounts.yaml").exists()


def write_accounts_yaml(
    skill_root: Path = SKILL_ROOT_DEFAULT,
    cny: float = 10_000_000,
    hkd: float = 10_000_000,
    usd: float = 50_000_000,
) -> bool:
    """Write config/accounts.yaml with the given initial cash amounts.

    Returns True if the file was written, False if it already existed (skip).
    Call this BEFORE lifecycle.start() to set custom initial cash.

    Initial cash is a one-time setting: after the service runs for the first
    time, accounts are seeded from this file into SQLite. Subsequent edits to
    accounts.yaml are ignored — changing amounts requires rebuilding the
    account, which loses all trade history.
    """
    _ensure_dirs(skill_root)
    target = skill_root / "config" / "accounts.yaml"
    if target.exists():
        return False
    content = (
        "accounts:\n"
        "  acc_default:\n"
        '    name: "default-paper-account"\n'
        "    base_currency: CNY\n"
        "    initial_cash:\n"
        f"      CNY: {int(cny)}\n"
        f"      HKD: {int(hkd)}\n"
        f"      USD: {int(usd)}\n"
        "    allowed_markets: [CN, HK, US]\n"
    )
    target.write_text(content)
    return True


def materialize_configs(skill_root: Path = SKILL_ROOT_DEFAULT) -> None:
    """Copy default configs into `skill_root/config/`, skipping any file
    that already exists (so user edits survive skill upgrades).

    Quote upstream credentials are NOT written into yaml — set env vars per
    ``auth`` in ``config/main.yaml`` (``STOCKI_*`` for bearer, ``STOCKI_GATEWAY_URL`` for none;
    see ``INSTALL.md``).
    """
    _ensure_dirs(skill_root)
    cfg_dir = skill_root / "config"
    tpl_dir = _templates_dir()
    for name in CONFIG_FILES:
        target = cfg_dir / name
        if target.exists():
            continue  # preserve user edits
        src = tpl_dir / name
        body = src.read_text(encoding="utf-8")
        target.write_text(body, encoding="utf-8")


def ensure_token(skill_root: Path = SKILL_ROOT_DEFAULT) -> str:
    """Generate `state/token` (32-byte hex, mode 0600) if missing.

    Returns the token string. Idempotent: re-running returns the same
    token. The file is mode 0600 so other local users can't read it.
    """
    _ensure_dirs(skill_root)
    tok = skill_root / "state" / "token"
    if tok.exists():
        existing = tok.read_text().strip()
        if existing:
            return existing
    token = secrets.token_hex(32)
    tok.write_text(token)
    os.chmod(tok, 0o600)
    return token


def ensure_venv(skill_root: Path = SKILL_ROOT_DEFAULT) -> Path:
    """Create `.venv/` and pip-install the skill package.

    Returns the path to the venv's python interpreter. This step is
    network-bound on first call (pip resolves trading-assistant's deps)
    and should be expected to take 30-180 s.
    """
    venv = skill_root / ".venv"
    if not venv_has_python(venv):
        subprocess.run(
            [sys.executable, "-m", "venv", str(venv)],
            check=True,
        )
    pip = venv_pip(venv)
    subprocess.run(
        [str(pip), "install", "--quiet", "-e", str(skill_root)],
        check=True,
    )
    return venv_python(venv)


def run(skill_root: Path = SKILL_ROOT_DEFAULT) -> None:
    """Convenience: do all three steps. Called by lifecycle.start() on first run."""
    materialize_configs(skill_root)
    ensure_token(skill_root)
    ensure_venv(skill_root)


if __name__ == "__main__":
    run()