support register custom feature ops easily

qlib/data/ops.py

@@ -13,6 +13,7 @@ from scipy.stats import percentileofscore
 
 from .base import Expression, ExpressionOps
 from ..log import get_module_logger
+from ..utils import OpsWrapper
 
 try:
     from ._libs.rolling import rolling_slope, rolling_rsquare, rolling_resi
@@ -21,53 +22,6 @@ except ImportError as err:
     print("Do not import qlib package in the repository directory!")
     raise
 
-__all__ = (
-    "Ref",
-    "Max",
-    "Min",
-    "Sum",
-    "Mean",
-    "Std",
-    "Var",
-    "Skew",
-    "Kurt",
-    "Med",
-    "Mad",
-    "Slope",
-    "Rsquare",
-    "Resi",
-    "Rank",
-    "Quantile",
-    "Count",
-    "EMA",
-    "WMA",
-    "Corr",
-    "Cov",
-    "Delta",
-    "Abs",
-    "Sign",
-    "Log",
-    "Power",
-    "Add",
-    "Sub",
-    "Mul",
-    "Div",
-    "Greater",
-    "Less",
-    "And",
-    "Or",
-    "Not",
-    "Gt",
-    "Ge",
-    "Lt",
-    "Le",
-    "Eq",
-    "Ne",
-    "Mask",
-    "IdxMax",
-    "IdxMin",
-    "If",
-)
 
 np.seterr(invalid="ignore")
 
@@ -1430,3 +1384,55 @@ class Cov(PairRolling):
 
     def __init__(self, feature_left, feature_right, N):
         super(Cov, self).__init__(feature_left, feature_right, N, "cov")
+
+Operators = OpsWrapper()
+
+OpsList = [
+    Ref,
+    Max,
+    Min,
+    Sum,
+    Mean,
+    Std,
+    Var,
+    Skew,
+    Kurt,
+    Med,
+    Mad,
+    Slope,
+    Rsquare,
+    Resi,
+    Rank,
+    Quantile,
+    Count,
+    EMA,
+    WMA,
+    Corr,
+    Cov,
+    Delta,
+    Abs,
+    Sign,
+    Log,
+    Power,
+    Add,
+    Sub,
+    Mul,
+    Div,
+    Greater,
+    Less,
+    And,
+    Or,
+    Not,
+    Gt,
+    Ge,
+    Lt,
+    Le,
+    Eq,
+    Ne,
+    Mask,
+    IdxMax,
+    IdxMin,
+    If,
+]
+
+Operators.register(OpsList)

qlib/utils/__init__.py

@@ -162,7 +162,7 @@ def parse_field(field):
     # - $open+$close -> Feature("open")+Feature("close")
     if not isinstance(field, str):
         field = str(field)
-    return re.sub(r"\$(\w+)", r'Feature("\1")', field)
+    return re.sub(r"\$(\w+)", r'Feature("\1")', re.sub(r"(\w+\s*)\(", r"Operators.\1(", field))
 
 
 def get_module_by_module_path(module_path):
@@ -728,3 +728,20 @@ def load_dataset(path_or_obj):
     elif extension == ".csv":
         return pd.read_csv(path_or_obj, parse_dates=True, index_col=[0, 1])
     raise ValueError(f"unsupported file type `{extension}`")
+
+#################### Operator Wrapper #####################
+
+class OpsWrapper(object):
+    """Ops Wrapper"""
+
+    def __init__(self):
+        self._ops = {}
+    def register(self, ops_list):
+
+        for operator in ops_list:
+            self._ops[operator.__name__] = operator
+
+    def __getattr__(self, key):
+        if self._ops is {}:
+            raise AttributeError("Please run qlib.init() first using qlib to register ops")
+        return self._ops[key]