feat(factor): add Factor base class and registry

This commit is contained in:
GitHub
2026-06-12 19:43:54 +08:00
parent 7d6607b0cf
commit 67d9963f20
3 changed files with 177 additions and 0 deletions
+6
View File
@@ -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"]
+54
View File
@@ -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
+117
View File
@@ -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