mirror of
https://ghfast.top/https://github.com/aeroxw/easy_tdx_max.git
synced 2026-09-12 16:54:20 +08:00
feat(factor): add Factor base class and registry
This commit is contained in:
@@ -0,0 +1,6 @@
|
||||
# src/easy_tdx/factor/__init__.py
|
||||
"""因子研究模块。"""
|
||||
|
||||
from easy_tdx.factor.base import FACTORY_REGISTRY, Factor, register_factor
|
||||
|
||||
__all__ = ["Factor", "register_factor", "FACTORY_REGISTRY"]
|
||||
@@ -0,0 +1,54 @@
|
||||
# src/easy_tdx/factor/base.py
|
||||
"""因子基类与全局注册表。"""
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
import pandas as pd
|
||||
|
||||
|
||||
class Factor(ABC):
|
||||
"""因子基类 — 所有因子的抽象契约。
|
||||
|
||||
子类必须定义:
|
||||
name: str — 唯一标识,如 "momentum_20d"
|
||||
category: str — 分类:momentum / value / quality / volatility / technical / chanlun
|
||||
description: str — 人类可读描述
|
||||
inputs: tuple[str, ...] — 需要的列名,如 ("close", "vol")
|
||||
|
||||
并实现 compute(df) -> pd.Series。
|
||||
"""
|
||||
|
||||
name: str
|
||||
category: str
|
||||
description: str
|
||||
inputs: tuple[str, ...]
|
||||
|
||||
@abstractmethod
|
||||
def compute(self, df: pd.DataFrame) -> pd.Series:
|
||||
"""接收 OHLCV DataFrame,返回因子值序列(与 df 等长)。"""
|
||||
...
|
||||
|
||||
def __init__(self) -> None:
|
||||
for attr in ("name", "category", "description", "inputs"):
|
||||
if not hasattr(self, attr):
|
||||
raise TypeError(
|
||||
f"Factor 子类 {type(self).__name__} 必须定义类属性 '{attr}'"
|
||||
)
|
||||
|
||||
|
||||
FACTORY_REGISTRY: dict[str, type[Factor]] = {}
|
||||
|
||||
|
||||
def register_factor(cls: type[Factor]) -> type[Factor]:
|
||||
"""类装饰器,将 Factor 子类注册到全局表。
|
||||
|
||||
Raises:
|
||||
ValueError: 如果 name 已被注册。
|
||||
"""
|
||||
if cls.name in FACTORY_REGISTRY:
|
||||
raise ValueError(
|
||||
f"因子 '{cls.name}' 已注册(类: {FACTORY_REGISTRY[cls.name].__name__})"
|
||||
)
|
||||
FACTORY_REGISTRY[cls.name] = cls
|
||||
return cls
|
||||
@@ -0,0 +1,117 @@
|
||||
# tests/unit/test_factor_base.py
|
||||
"""Test Factor base class and registry."""
|
||||
from __future__ import annotations
|
||||
|
||||
import pandas as pd
|
||||
import pytest
|
||||
|
||||
from easy_tdx.factor.base import (
|
||||
FACTORY_REGISTRY,
|
||||
Factor,
|
||||
register_factor,
|
||||
)
|
||||
|
||||
|
||||
class _StubFactor(Factor):
|
||||
"""测试用因子。"""
|
||||
|
||||
name = "test_stub"
|
||||
category = "test"
|
||||
description = "stub for testing"
|
||||
inputs = ("close",)
|
||||
|
||||
def compute(self, df: pd.DataFrame) -> pd.Series:
|
||||
return df["close"].pct_change(1)
|
||||
|
||||
|
||||
class TestFactorABC:
|
||||
def test_cannot_instantiate_abc(self):
|
||||
with pytest.raises(TypeError):
|
||||
Factor() # type: ignore[abstract]
|
||||
|
||||
def test_subclass_must_define_name(self):
|
||||
class NoName(Factor):
|
||||
category = "test"
|
||||
description = "x"
|
||||
inputs = ("close",)
|
||||
|
||||
def compute(self, df):
|
||||
return df["close"]
|
||||
|
||||
with pytest.raises(TypeError):
|
||||
NoName()
|
||||
|
||||
def test_subclass_must_define_category(self):
|
||||
class NoCategory(Factor):
|
||||
name = "x"
|
||||
description = "x"
|
||||
inputs = ("close",)
|
||||
|
||||
def compute(self, df):
|
||||
return df["close"]
|
||||
|
||||
with pytest.raises(TypeError):
|
||||
NoCategory()
|
||||
|
||||
def test_subclass_must_implement_compute(self):
|
||||
class NoCompute(Factor):
|
||||
name = "x"
|
||||
category = "test"
|
||||
description = "x"
|
||||
inputs = ("close",)
|
||||
|
||||
with pytest.raises(TypeError):
|
||||
NoCompute()
|
||||
|
||||
def test_concrete_subclass_works(self):
|
||||
f = _StubFactor()
|
||||
assert f.name == "test_stub"
|
||||
assert f.category == "test"
|
||||
assert f.inputs == ("close",)
|
||||
|
||||
|
||||
class TestRegistry:
|
||||
def test_register_factor_decorator(self):
|
||||
@register_factor
|
||||
class RegFactor(Factor):
|
||||
name = "reg_test_factor"
|
||||
category = "test"
|
||||
description = "registered factor"
|
||||
inputs = ("close",)
|
||||
|
||||
def compute(self, df):
|
||||
return df["close"]
|
||||
|
||||
assert "reg_test_factor" in FACTORY_REGISTRY
|
||||
assert FACTORY_REGISTRY["reg_test_factor"] is RegFactor
|
||||
|
||||
def test_duplicate_name_raises(self):
|
||||
@register_factor
|
||||
class Dup(Factor):
|
||||
name = "dup_test_factor"
|
||||
category = "test"
|
||||
description = "dup"
|
||||
inputs = ("close",)
|
||||
|
||||
def compute(self, df):
|
||||
return df["close"]
|
||||
|
||||
with pytest.raises(ValueError, match="已注册"):
|
||||
@register_factor
|
||||
class Dup2(Factor):
|
||||
name = "dup_test_factor"
|
||||
category = "test"
|
||||
description = "dup2"
|
||||
inputs = ("close",)
|
||||
|
||||
def compute(self, df):
|
||||
return df["close"]
|
||||
|
||||
|
||||
class TestFactorCompute:
|
||||
def test_compute_returns_series(self):
|
||||
f = _StubFactor()
|
||||
df = pd.DataFrame({"close": [10.0, 11.0, 10.5, 12.0]})
|
||||
result = f.compute(df)
|
||||
assert isinstance(result, pd.Series)
|
||||
assert len(result) == 4
|
||||
Reference in New Issue
Block a user