mirror of
https://github.com/shy3130/tick-stock-panel.git
synced 2026-11-02 10:26:13 +00:00
自定义因子新增第二档编写语言 (lang=python): 受限 Python 表达式块,
只返回单个 Polars 表达式。三档编译层就此齐备: T1 DSL < T2 表达式级
Python < T3 自由 Python 策略。
- app/factors/py_expr.py: AST 白名单文法 (仅赋值/表达式语句; 名字限
ctx/pl 与已赋值局部量; 禁 import/def/lambda/循环/推导式/下划线属性/
map_*/apply/IO 类调用) + {ctx, pl} + 空 builtins 求值; ctx.col 构造期
登记依赖 (无需推导); warmup 从窗口常量保守求和 (显式声明只能增大,
少算失真是错误方向)。
- 贯通链路: FactorSpec.lang (默认 dsl, 存量零影响) → to_spec 编译分支 →
materialize 按 lang 分发 (T2 结果同样进 P2b 内容寻址缓存, 依赖即缓存键
输入) → API /validate 与 /trial 支持 lang=python。
- 验收 (docs/ext-architecture-plan.md P3): 看板口径 SKDJ — 9日RSV
(分母 0 时取 50) → EMA(α=0.5) 两级 — 的 T2 单表达式与策略文件的分阶段
列实现逐元素一致; 19 个契约测试 (含依赖登记/warmup 下限/11 种文法拒绝
形态/物化与缓存贯通), 相关回归 568 全过, ruff 新增文件全清。
Signed-off-by: shy3130 <shy3130@qq.com>
182 lines
7.5 KiB
Python
182 lines
7.5 KiB
Python
"""T2 表达式级 Python (P3) 契约测试。
|
||
|
||
验收 (docs/ext-architecture-plan.md P3): SKDJ 类口径零 DSL 可表达 —
|
||
看板口径 SKDJ (9日RSV → EMA(α=0.5) 两级, 分母为 0 时 RSV 取 50) 与
|
||
策略文件里的分阶段 Polars 实现**逐元素一致**。
|
||
|
||
其余契约:
|
||
- ctx.col 构造期登记依赖 (进 FactorSpec.dependencies → 计划/缓存同源);
|
||
- warmup 从窗口常量保守推断 (只多不少);
|
||
- 文法白名单 fail-closed (import/def/lambda/循环/推导式/下划线属性/
|
||
回调类方法/IO 类调用/未定义名字/非表达式结果全部拒绝);
|
||
- 存储与物化链路贯通 (to_spec 注册 → materialize 计算 → 内容寻址缓存)。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import numpy as np
|
||
import polars as pl
|
||
import pytest
|
||
|
||
from app.config import settings
|
||
from app.factors import materialize_cache
|
||
from app.factors.py_expr import compile_python_expression
|
||
from app.factors.registry import unregister_factor
|
||
from app.factors.store import register_definition
|
||
|
||
# 看板口径 SKDJ — 与 custom_v1_kanban_strict 策略文件逐行同源 (分阶段列实现)
|
||
SKDJ_T2 = '''
|
||
llv = ctx.col("low").rolling_min(window_size=9).over("symbol")
|
||
hhv = ctx.col("high").rolling_max(window_size=9).over("symbol")
|
||
rsv = pl.when((hhv - llv) > 0).then((ctx.col("close") - llv) / (hhv - llv) * 100.0).otherwise(pl.lit(50.0))
|
||
k = rsv.ewm_mean(alpha=0.5, adjust=False).over("symbol")
|
||
d = k.ewm_mean(alpha=0.5, adjust=False).over("symbol")
|
||
d
|
||
'''
|
||
|
||
|
||
def _skdj_reference(df: pl.DataFrame) -> pl.Series:
|
||
d = df.sort(["symbol", "date"]).with_columns([
|
||
pl.col("low").rolling_min(window_size=9).over("symbol").alias("_llv9"),
|
||
pl.col("high").rolling_max(window_size=9).over("symbol").alias("_hhv9"),
|
||
])
|
||
denom = pl.col("_hhv9") - pl.col("_llv9")
|
||
d = d.with_columns(
|
||
pl.when(denom > 0)
|
||
.then((pl.col("close") - pl.col("_llv9")) / denom * 100.0)
|
||
.otherwise(pl.lit(50.0))
|
||
.alias("_rsv")
|
||
)
|
||
d = d.with_columns(pl.col("_rsv").ewm_mean(alpha=0.5, adjust=False).over("symbol").alias("_sk_k"))
|
||
d = d.with_columns(pl.col("_sk_k").ewm_mean(alpha=0.5, adjust=False).over("symbol").alias("_sk_d"))
|
||
return d["_sk_d"]
|
||
|
||
|
||
def _frame(n_days: int = 120) -> pl.DataFrame:
|
||
rng = np.random.default_rng(11)
|
||
rows = []
|
||
for sym in ("A", "B", "C"):
|
||
base = rng.uniform(8, 12)
|
||
for i in range(n_days):
|
||
close = base + rng.normal(0, 0.3) + i * 0.005
|
||
rows.append({
|
||
"symbol": sym, "date": 20260101 + i,
|
||
"open": close * 0.999, "high": close + 0.2, "low": close - 0.2, "close": close,
|
||
"volume": rng.uniform(1e4, 1e6),
|
||
})
|
||
return pl.DataFrame(rows).sort(["symbol", "date"])
|
||
|
||
|
||
# ── 验收: SKDJ 等价 ──────────────────────────────────────
|
||
|
||
def test_skdj_t2_matches_staged_reference_exactly() -> None:
|
||
compiled = compile_python_expression(SKDJ_T2)
|
||
assert compiled.ok, [e.message for e in compiled.errors]
|
||
frame = _frame()
|
||
via_t2 = frame.select(compiled.expr.alias("d"))["d"]
|
||
reference = _skdj_reference(frame)
|
||
assert via_t2.equals(reference)
|
||
# 窗口预热区 (前 8 行 RSV 初始 50 口径) 也参与比较 — equals 全量逐元素
|
||
|
||
|
||
def test_skdj_dependencies_and_warmup() -> None:
|
||
compiled = compile_python_expression(SKDJ_T2)
|
||
assert compiled.dependencies == frozenset({"low", "high", "close"})
|
||
# rolling 9×2 (low/high 各一) + ewm(α=0.5→span3) ×2, +1 → 25 (保守, 只多不少)
|
||
assert compiled.warmup_bars == 25
|
||
|
||
|
||
# ── 编译契约 ─────────────────────────────────────────────
|
||
|
||
def test_dependencies_tracked_from_ctx_col() -> None:
|
||
compiled = compile_python_expression('ctx.col("close") / ctx.col("ma20") - 1')
|
||
assert compiled.ok
|
||
assert compiled.dependencies == frozenset({"close", "ma20"})
|
||
|
||
|
||
def test_declared_warmup_cannot_shrink_below_inferred() -> None:
|
||
register_definition({
|
||
"id": "uf_py_w", "label": "w", "kind": "custom", "lang": "python",
|
||
"formula": 'ctx.col("close").rolling_mean(window_size=30).over("symbol")',
|
||
"warmup_bars": 5, "status": "active", # 声明小于推断 → 取推断 (下限)
|
||
})
|
||
from app.factors.registry import get_factor
|
||
try:
|
||
assert get_factor("uf_py_w").warmup_bars == 31
|
||
finally:
|
||
unregister_factor("uf_py_w")
|
||
|
||
|
||
@pytest.mark.parametrize("source, fragment", [
|
||
("import os", "不允许的语法: Import"),
|
||
("def f():\n return 1", "不允许的语法"),
|
||
("ctx.col('close').map_batches(lambda s: s)", "不允许的语法: Lambda"),
|
||
("x = 0\nfor i in range(3):\n x = x + i", "不允许的语法: For"),
|
||
("y = [ctx.col('close') for _ in range(2)]", "不允许的语法: ListComp"),
|
||
("ctx.col('close').__class__", "下划线属性"),
|
||
("ctx.col('close').map_elements(str)", "回调"),
|
||
("pl.read_csv('/etc/passwd')", "IO 类调用"),
|
||
("unknown_name + 1", "未定义的名字"),
|
||
("42", "结果必须是单个 Polars 表达式"),
|
||
("pl.col('close') + 1", "至少通过 ctx.col"), # 绕过 ctx → 无依赖登记
|
||
])
|
||
def test_grammar_rejects_unsafe_or_malformed(source: str, fragment: str) -> None:
|
||
compiled = compile_python_expression(source)
|
||
assert not compiled.ok
|
||
messages = [e.message for e in compiled.errors]
|
||
assert any(fragment in m for m in messages), messages
|
||
|
||
|
||
def test_builtin_functions_unavailable() -> None:
|
||
compiled = compile_python_expression("len(ctx.col('close'))")
|
||
assert not compiled.ok # 空 builtins: NameError → E212 求值失败
|
||
|
||
|
||
def test_declared_warmup_can_grow() -> None:
|
||
register_definition({
|
||
"id": "uf_py_w2", "label": "w2", "kind": "custom", "lang": "python",
|
||
"formula": 'ctx.col("close").rolling_mean(window_size=10).over("symbol")',
|
||
"warmup_bars": 60, "status": "active",
|
||
})
|
||
from app.factors.registry import get_factor
|
||
try:
|
||
assert get_factor("uf_py_w2").warmup_bars == 60 # 声明更大 → 生效
|
||
finally:
|
||
unregister_factor("uf_py_w2")
|
||
|
||
|
||
def test_unknown_lang_rejected() -> None:
|
||
with pytest.raises(ValueError, match="未知 lang"):
|
||
register_definition({
|
||
"id": "uf_py_bad", "label": "b", "kind": "custom", "lang": "sql",
|
||
"formula": "select 1", "status": "active",
|
||
})
|
||
|
||
|
||
# ── 物化链路贯通 ─────────────────────────────────────────
|
||
|
||
@pytest.fixture(autouse=True)
|
||
def _isolated_cache(tmp_path, monkeypatch):
|
||
monkeypatch.setattr(settings, "data_dir", tmp_path)
|
||
materialize_cache.invalidate(disk=True)
|
||
yield
|
||
materialize_cache.invalidate(disk=True)
|
||
|
||
|
||
def test_python_factor_materializes_and_caches() -> None:
|
||
from app.strategy.scoring import materialize_scoring_columns
|
||
|
||
register_definition({
|
||
"id": "uf_py_k", "label": "SKDJ-K", "kind": "custom", "lang": "python",
|
||
"formula": SKDJ_T2, "status": "active",
|
||
})
|
||
frame = _frame(120) # 小帧走缓存旁路 (语义不变), 物化链路本身不受影响
|
||
try:
|
||
result = materialize_scoring_columns(frame, ["uf_py_k"])
|
||
assert "uf_py_k" in result.columns
|
||
assert result["uf_py_k"].equals(_skdj_reference(frame))
|
||
# 重复物化 (帧上无该列) → 结果一致
|
||
again = materialize_scoring_columns(frame, ["uf_py_k"])
|
||
assert again["uf_py_k"].equals(result["uf_py_k"])
|
||
finally:
|
||
unregister_factor("uf_py_k")
|