Add minute bar feature pipeline

This commit is contained in:
Yuxuan Yan
2026-06-16 13:57:17 +08:00
parent 17fa75495d
commit 83a006bbe4
19 changed files with 1289 additions and 11 deletions
+19 -5
View File
@@ -5,7 +5,7 @@ position weights. Subclasses implement :meth:`signal` — the raw, unnormalized
score. The base class turns a signal into cross-sectionally z-scored weights
via :meth:`to_weights` (override it for a different normalization).
"""
from abc import ABC, abstractmethod
from abc import ABC
import numpy as np
import pandas as pd
@@ -15,15 +15,14 @@ class BaseAlpha(ABC):
"""A position-weight alpha over a cross-section of stocks.
Concrete subclasses must set a unique class-level :attr:`name` (the registry
key) and implement :meth:`signal`. Construct subclasses with their own typed
parameters (e.g. ``lookback``); the factory passes only the parameters a
given ``__init__`` accepts.
key) and implement either :meth:`signal` or :meth:`signal_from_data`.
Construct subclasses with their own typed parameters (e.g. ``lookback``);
the factory passes only the parameters a given ``__init__`` accepts.
"""
#: Unique registry key. Every concrete alpha must set this to a non-empty str.
name: str = ""
@abstractmethod
def signal(self, close: pd.DataFrame) -> pd.DataFrame:
"""Compute the raw signal.
@@ -34,6 +33,21 @@ class BaseAlpha(ABC):
A wide DataFrame aligned to ``close`` where higher values indicate a
stronger long. Use NaN where the signal is undefined.
"""
raise NotImplementedError(
f"{type(self).__name__} must implement signal() or signal_from_data()"
)
def signal_from_data(
self,
data: pd.DataFrame,
close: pd.DataFrame,
) -> pd.DataFrame:
"""Compute the raw signal from long daily data plus wide closes.
Feature-aware alphas can override this to pivot joined feature columns
from ``data``. The default preserves the existing close-only alpha API.
"""
return self.signal(close)
def to_weights(self, signal: pd.DataFrame) -> pd.DataFrame:
"""Cross-sectionally z-score a signal into signed position weights.
+6 -1
View File
@@ -56,6 +56,10 @@ def list_(alpha_modules):
@click.option("--output-dir", default="alphas", help="Directory to save alpha parquet")
@click.option("--lookback", default=5, type=int, help="Lookback days")
@click.option("--vol-window", default=20, type=int, help="Volatility window (reversal_vol only)")
@click.option(
"--feature-path", "feature_paths", multiple=True,
help="Daily feature parquet file/dataset to left-join on symbol_id,date (repeatable)",
)
@click.option(
"--alpha-module", "alpha_modules", multiple=True,
help="External module(s) to import so their alphas register (dotted path or .py file)",
@@ -74,7 +78,7 @@ def list_(alpha_modules):
help="Most-liquid names kept per date when --liquid-universe is set",
)
def compute(data_path, alpha_name, alpha_type, output_dir, lookback, vol_window,
alpha_modules, extra_params, liquid_universe, universe_top_n):
feature_paths, alpha_modules, extra_params, liquid_universe, universe_top_n):
"""Compute one alpha from raw data and save as parquet."""
for spec in alpha_modules:
load_alpha_module(spec)
@@ -100,6 +104,7 @@ def compute(data_path, alpha_name, alpha_type, output_dir, lookback, vol_window,
alpha_name=alpha_name,
alpha_type=alpha_type,
universe=universe,
feature_paths=feature_paths,
**params,
)
+51 -2
View File
@@ -7,12 +7,15 @@ through :mod:`pipeline.alpha.registry`.
"""
import logging
from pathlib import Path
from typing import Iterable
import numpy as np
import pandas as pd
from pipeline.alpha.registry import get_alpha
from pipeline.common.schema import ALPHA_COLUMNS
from pipeline.features.compute import FEATURE_KEY_COLUMNS, read_feature_frames, validate_feature_frame
logger = logging.getLogger(__name__)
@@ -33,6 +36,38 @@ def _pivot_open(df: pd.DataFrame) -> pd.DataFrame:
return pivot.sort_index()
def join_feature_frames(
data: pd.DataFrame,
feature_frames: Iterable[pd.DataFrame],
) -> pd.DataFrame:
"""Left-join validated daily feature frames onto long daily data."""
out = data.copy()
out["date"] = pd.to_datetime(out["date"])
existing = set(out.columns)
joined_cols: list[str] = []
for frame in feature_frames:
features = validate_feature_frame(frame)
feature_cols = [col for col in features.columns if col not in FEATURE_KEY_COLUMNS]
overlap = sorted(existing.intersection(feature_cols))
if overlap:
raise ValueError(
f"Feature columns conflict with existing daily data columns: {overlap}"
)
out = out.merge(
features,
on=FEATURE_KEY_COLUMNS,
how="left",
validate="many_to_one",
)
existing.update(feature_cols)
joined_cols.extend(feature_cols)
if joined_cols:
logger.info("Joined feature columns into daily data: %s", joined_cols)
return out
def _forward_open_to_open_returns(open_: pd.DataFrame) -> pd.DataFrame:
"""Return earned by a close-formed signal after next-open execution.
@@ -105,6 +140,8 @@ def compute_alpha(
alpha_name: str,
alpha_type: str,
universe: dict | None = None,
feature_paths: Iterable[str | Path] | None = None,
feature_frames: Iterable[pd.DataFrame] | None = None,
**params,
) -> pd.DataFrame:
"""Compute alpha weights from raw data.
@@ -118,6 +155,10 @@ def compute_alpha(
:func:`investable_universe_mask`) *before* it is turned into
weights, so unheld names get weight 0. Keys are forwarded as keyword
arguments to :func:`investable_universe_mask`.
feature_paths: Optional parquet files/datasets keyed by ``symbol_id``
and ``date``. Their numeric feature columns are left-joined onto
``data`` before alpha logic runs.
feature_frames: Optional in-memory feature frames with the same schema.
**params: Constructor parameters for the alpha (e.g. ``lookback``,
``vol_window``). Only the params the alpha's ``__init__`` accepts are
used; extras are ignored.
@@ -128,12 +169,20 @@ def compute_alpha(
Raises:
KeyError: If ``alpha_type`` is not registered.
"""
feature_inputs: list[pd.DataFrame] = []
if feature_paths:
feature_inputs.extend(read_feature_frames(feature_paths))
if feature_frames:
feature_inputs.extend(feature_frames)
if feature_inputs:
data = join_feature_frames(data, feature_inputs)
alpha = get_alpha(alpha_type, **params)
close = _pivot_close(data)
signal = alpha.signal_from_data(data, close)
if universe is None:
weights = alpha.weights(close)
weights = alpha.to_weights(signal)
else:
signal = alpha.signal(close)
mask = investable_universe_mask(data, signal, **universe)
weights = alpha.to_weights(signal.where(mask))
+18
View File
@@ -26,6 +26,24 @@ DATA_COLUMNS: Final[list[str]] = [
"pcfNcfTTM", # float64: P/CF (net cash flow, TTM)
]
# Required columns for raw intraday minute bar parquet files.
MINUTE_BAR_COLUMNS: Final[list[str]] = [
"symbol_id", # str: internal code like 'sh600000'
"symbol_name", # str: stock name like '浦发银行'
"datetime", # datetime64: intraday bar timestamp
"date", # date component, aligned with daily DATA_COLUMNS date
"time", # str: HH:MM:SS bar time
"frequency", # str: e.g. '5m'
"open", # float64
"high", # float64
"low", # float64
"close", # float64
"volume", # float64 (shares)
"amount", # float64 (turnover in yuan, raw/unadjusted)
"vwap", # float64: amount / volume
"adjustflag", # str: baostock adjustment flag; '3' for raw/unadjusted
]
# Required columns for alpha parquet files.
# Alphas are position WEIGHTS: positive=long, negative=short.
ALPHA_COLUMNS: Final[list[str]] = [
+35 -1
View File
@@ -3,7 +3,7 @@
import click
from datetime import date
from pipeline.data.downloader import download_universe
from pipeline.data.downloader import download_minute_universe, download_universe
@click.group(name="data")
@@ -42,3 +42,37 @@ def download(universe, start_date, end_date, output_dir, symbols, chunk_size, ad
f"{stats['n_rows']:,} bars, {stats['date_min']}{stats['date_max']}"
)
click.echo(f"Dataset: {stats['dataset_path']}")
@data.command("download-minute")
@click.option(
"--universe", default="csi500",
help="Which universe: hs300, csi500, all (~5000 A-shares), file path, or comma-separated symbols",
)
@click.option("--start-date", default="2017-01-01", help="Start date YYYY-MM-DD")
@click.option("--end-date", default=str(date.today()), help="End date YYYY-MM-DD")
@click.option("--output-dir", default="data/minute_bars", help="Root for the partitioned dataset")
@click.option("--symbols", default=0, type=int, help="Max symbols (0=all)")
@click.option("--chunk-size", default=100, type=int, help="Symbols per durability flush")
@click.option("--frequency", default="5", help="Minute frequency: 5, 15, 30, or 60")
def download_minute(universe, start_date, end_date, output_dir, symbols, chunk_size, frequency):
"""Download raw Baostock minute bars into a partitioned parquet dataset.
Writes ``{output_dir}/{universe}/frequency=5m/month=YYYY-MM/*.pq`` for the
default 5-minute frequency.
"""
stats = download_minute_universe(
universe=universe,
start_date=start_date,
end_date=end_date,
output_dir=output_dir,
max_symbols=symbols,
chunk_size=chunk_size,
frequency=frequency,
)
click.echo(
f"\nSummary: {stats['n_symbols']}/{stats['n_requested']} symbols, "
f"{stats['n_rows']:,} bars, {stats['date_min']}{stats['date_max']}, "
f"frequency={stats['frequency']}"
)
click.echo(f"Dataset: {stats['dataset_path']}")
+126 -2
View File
@@ -11,9 +11,13 @@ import pyarrow.dataset as pads
# Reuse existing downloader and universe modules
sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
from data.downloader import download_daily_batch
from data.downloader import (
_normalize_minute_frequency,
download_daily_batch,
download_minute_batch,
)
from data.universe import get_all_stocks, get_hs300_stocks, get_zz500_stocks
from pipeline.common.schema import DATA_COLUMNS
from pipeline.common.schema import DATA_COLUMNS, MINUTE_BAR_COLUMNS
logger = logging.getLogger(__name__)
@@ -89,6 +93,25 @@ def _write_month_partitions(df: pd.DataFrame, base_dir: Path, basename_prefix: s
)
def _write_minute_partitions(df: pd.DataFrame, base_dir: Path, basename_prefix: str) -> None:
"""Append rows to a Hive-partitioned minute dataset.
Layout: ``frequency=5m/month=YYYY-MM/*.pq``.
"""
out = df.copy()
out["month"] = pd.to_datetime(out["date"]).dt.strftime("%Y-%m")
table = pa.Table.from_pandas(out, preserve_index=False)
pads.write_dataset(
table,
str(base_dir),
format="parquet",
partitioning=["frequency", "month"],
partitioning_flavor="hive",
basename_template=f"{basename_prefix}-{{i}}.pq",
existing_data_behavior="overwrite_or_ignore",
)
def download_universe(
universe: str = "csi500",
start_date: str = "2017-01-01",
@@ -177,3 +200,104 @@ def download_universe(
"date_min": None if date_min is None else str(pd.Timestamp(date_min).date()),
"date_max": None if date_max is None else str(pd.Timestamp(date_max).date()),
}
def download_minute_universe(
universe: str = "csi500",
start_date: str = "2017-01-01",
end_date: str = "2026-12-31",
output_dir: str = "data/minute_bars",
max_symbols: int = 0,
chunk_size: int = 100,
frequency: str | int = 5,
) -> dict:
"""Download raw minute bars into a frequency/month-partitioned dataset.
Args:
universe: ``hs300``, ``csi500``, ``all``/``full``, a file path, or a
comma-separated symbol list.
start_date, end_date: ``YYYY-MM-DD`` bounds.
output_dir: Root under which ``{universe}/frequency=5m/month=YYYY-MM``
is written.
max_symbols: Cap on symbols (0 = all).
chunk_size: Symbols per durability flush.
frequency: Minute interval. ``5``/``"5"``/``"5m"`` are 5-minute bars.
Returns:
Stats dict with dataset path, row count, symbol count, date range, and
frequency label.
"""
_, frequency_label = _normalize_minute_frequency(frequency)
constituents = _resolve_universe(universe, max_symbols)
symbols = constituents["symbol_id"].tolist()
names = dict(zip(constituents["symbol_id"], constituents["symbol_name"]))
n_requested = len(symbols)
logger.info(
"Minute universe %s: %d symbols, %s%s, frequency=%s",
universe,
n_requested,
start_date,
end_date,
frequency,
)
base_dir = Path(output_dir) / universe
target_frequency_dir = base_dir / f"frequency={frequency_label}"
if target_frequency_dir.exists():
shutil.rmtree(target_frequency_dir)
base_dir.mkdir(parents=True, exist_ok=True)
buffer: list[pd.DataFrame] = []
chunk_idx = 0
succeeded = 0
n_rows = 0
date_min = None
date_max = None
def flush() -> None:
nonlocal buffer, chunk_idx, n_rows, date_min, date_max
if not buffer:
return
chunk = pd.concat(buffer, ignore_index=True)
_write_minute_partitions(chunk, base_dir, basename_prefix=f"chunk{chunk_idx:04d}")
n_rows += len(chunk)
cmin, cmax = chunk["date"].min(), chunk["date"].max()
date_min = cmin if date_min is None else min(date_min, cmin)
date_max = cmax if date_max is None else max(date_max, cmax)
logger.info(
"Flushed minute chunk %d: %d rows (%d symbols done)",
chunk_idx,
len(chunk),
succeeded,
)
buffer = []
chunk_idx += 1
for i, (symbol, df) in enumerate(
download_minute_batch(symbols, start_date, end_date, frequency=frequency), start=1
):
if df is None:
logger.warning(" %s: no minute data", symbol)
else:
df["symbol_id"] = symbol
df["symbol_name"] = names.get(symbol, symbol)
buffer.append(df[MINUTE_BAR_COLUMNS])
succeeded += 1
if len(buffer) >= chunk_size:
flush()
if i % 100 == 0:
logger.info("Minute progress: %d/%d symbols", i, n_requested)
flush()
if succeeded == 0:
raise RuntimeError("No minute data downloaded for any symbol")
return {
"dataset_path": str(base_dir),
"frequency": frequency_label,
"n_symbols": succeeded,
"n_requested": n_requested,
"n_rows": n_rows,
"date_min": None if date_min is None else str(pd.Timestamp(date_min).date()),
"date_max": None if date_max is None else str(pd.Timestamp(date_max).date()),
}
+1
View File
@@ -0,0 +1 @@
"""Daily feature plugin package."""
+34
View File
@@ -0,0 +1,34 @@
"""Base class for daily feature plugins."""
from abc import ABC, abstractmethod
import pandas as pd
class BaseFeature(ABC):
"""Aggregate raw minute bars into daily, symbol-keyed feature columns."""
#: Unique registry key. Every concrete feature must set this to a non-empty str.
name: str = ""
@abstractmethod
def compute(
self,
minute: pd.DataFrame,
daily: pd.DataFrame | None = None,
) -> pd.DataFrame:
"""Compute daily features.
Args:
minute: Raw minute bars with ``symbol_id`` and ``date`` keys.
daily: Optional daily data frame for calendar alignment or
reference daily columns.
Returns:
DataFrame with ``symbol_id``, ``date``, and one or more numeric
feature columns.
"""
def __repr__(self) -> str:
params = ", ".join(f"{k}={v!r}" for k, v in vars(self).items())
return f"{type(self).__name__}({params})"
+108
View File
@@ -0,0 +1,108 @@
"""CLI for daily feature computation."""
import os
import click
import pandas as pd
from pipeline.features.compute import compute_feature
from pipeline.features.registry import available_features, load_feature_module
@click.group(name="feature")
def feature():
"""Compute daily feature parquet files from minute bars."""
def _coerce(value: str):
"""Best-effort coercion of a CLI string to int, then float, else str."""
for cast in (int, float):
try:
return cast(value)
except ValueError:
continue
return value
def _parse_params(pairs: tuple[str, ...]) -> dict:
"""Parse repeated ``name=value`` options into a params dict."""
params: dict = {}
for pair in pairs:
if "=" not in pair:
raise click.BadParameter(f"--param must be name=value, got '{pair}'")
key, value = pair.split("=", 1)
params[key.strip()] = _coerce(value.strip())
return params
@feature.command("list")
@click.option(
"--feature-module", "feature_modules", multiple=True,
help="External module(s) to import first (dotted path or .py file)",
)
def list_(feature_modules):
"""List the registered feature types."""
for spec in feature_modules:
load_feature_module(spec)
for name in available_features():
click.echo(name)
@feature.command("compute")
@click.option("--minute-path", required=True, help="Path to minute parquet dataset/file")
@click.option("--daily-path", default=None, help="Optional daily data parquet for alignment")
@click.option("--feature-type", required=True, help="Registry key of the feature class")
@click.option("--feature-name", required=True, help="Name for this feature run/output file")
@click.option("--output-dir", default="features", help="Directory to save feature parquet")
@click.option(
"--feature-module", "feature_modules", multiple=True,
help="External module(s) to import so their features register (dotted path or .py file)",
)
@click.option(
"--param", "extra_params", multiple=True,
help="Extra feature constructor param as name=value (repeatable)",
)
def compute(
minute_path,
daily_path,
feature_type,
feature_name,
output_dir,
feature_modules,
extra_params,
):
"""Compute one daily feature file from raw minute bars."""
for spec in feature_modules:
load_feature_module(spec)
options = available_features()
if feature_type not in options:
raise click.BadParameter(
f"Unknown feature-type '{feature_type}'. Available: {options}. "
f"Use --feature-module to register an external feature.",
param_hint="--feature-type",
)
minute = pd.read_parquet(minute_path)
click.echo(f"Loaded minute bars: {len(minute):,} rows from {minute_path}")
daily = None
if daily_path:
daily = pd.read_parquet(daily_path)
click.echo(f"Loaded daily data: {len(daily):,} rows from {daily_path}")
result = compute_feature(
minute=minute,
daily=daily,
feature_type=feature_type,
**_parse_params(extra_params),
)
os.makedirs(output_dir, exist_ok=True)
out_path = f"{output_dir}/{feature_name}.pq"
result.to_parquet(out_path, index=False)
feature_cols = [col for col in result.columns if col not in ("symbol_id", "date")]
click.echo(
f"Saved feature: {out_path} ({len(result):,} rows, "
f"{len(feature_cols)} columns)"
)
+75
View File
@@ -0,0 +1,75 @@
"""Feature computation and validation."""
import logging
from pathlib import Path
from typing import Iterable
import pandas as pd
from pandas.api.types import is_numeric_dtype
from pipeline.features.registry import get_feature
logger = logging.getLogger(__name__)
FEATURE_KEY_COLUMNS = ["symbol_id", "date"]
def validate_feature_frame(features: pd.DataFrame) -> pd.DataFrame:
"""Validate and normalize a daily feature frame.
A valid feature frame is keyed by unique ``symbol_id,date`` rows and has at
least one numeric feature column beyond those keys.
"""
duplicated = features.columns[features.columns.duplicated()].tolist()
if duplicated:
raise ValueError(f"Feature output has duplicate columns: {duplicated}")
missing = [col for col in FEATURE_KEY_COLUMNS if col not in features.columns]
if missing:
raise ValueError(f"Feature output missing required columns: {missing}")
out = features.copy()
out["date"] = pd.to_datetime(out["date"])
if out.duplicated(FEATURE_KEY_COLUMNS).any():
raise ValueError("Feature output has duplicate symbol_id,date rows")
feature_cols = [col for col in out.columns if col not in FEATURE_KEY_COLUMNS]
if not feature_cols:
raise ValueError("Feature output must include at least one feature column")
non_numeric = [col for col in feature_cols if not is_numeric_dtype(out[col])]
if non_numeric:
raise ValueError(f"Feature columns must be numeric: {non_numeric}")
out = out[FEATURE_KEY_COLUMNS + feature_cols].copy()
return out.sort_values(FEATURE_KEY_COLUMNS).reset_index(drop=True)
def compute_feature(
minute: pd.DataFrame,
feature_type: str,
daily: pd.DataFrame | None = None,
**params,
) -> pd.DataFrame:
"""Compute one registered feature from raw minute bars."""
feature = get_feature(feature_type, **params)
result = validate_feature_frame(feature.compute(minute=minute, daily=daily))
feature_cols = [col for col in result.columns if col not in FEATURE_KEY_COLUMNS]
logger.info(
"Feature '%s' (%r): %d symbols × %d dates, columns=%s",
feature_type,
feature,
result["symbol_id"].nunique(),
result["date"].nunique(),
feature_cols,
)
return result
def read_feature_frames(feature_paths: Iterable[str | Path]) -> list[pd.DataFrame]:
"""Read and validate feature parquet files."""
return [
validate_feature_frame(pd.read_parquet(path))
for path in feature_paths
]
+3
View File
@@ -0,0 +1,3 @@
"""Built-in feature library."""
from pipeline.features.library import minute_daily_summary # noqa: F401
@@ -0,0 +1,84 @@
"""Daily summary features derived from raw minute bars."""
import numpy as np
import pandas as pd
from pipeline.features.base import BaseFeature
from pipeline.features.registry import register_feature
@register_feature
class MinuteDailySummaryFeature(BaseFeature):
"""Aggregate intraday bars into daily summary columns."""
name = "minute_daily_summary"
def compute(
self,
minute: pd.DataFrame,
daily: pd.DataFrame | None = None,
) -> pd.DataFrame:
minute = minute.copy()
minute["date"] = pd.to_datetime(minute["date"])
sort_cols = ["symbol_id", "date"]
if "datetime" in minute.columns:
minute["datetime"] = pd.to_datetime(minute["datetime"])
sort_cols.append("datetime")
elif "time" in minute.columns:
sort_cols.append("time")
minute = minute.sort_values(sort_cols)
grouped = minute.groupby(["symbol_id", "date"], sort=True)
summary = grouped.agg(
minute_bar_count=("close", "count"),
first_open=("open", "first"),
last_close=("close", "last"),
high=("high", "max"),
low=("low", "min"),
volume_sum=("volume", "sum"),
amount_sum=("amount", "sum"),
)
summary["minute_intraday_return"] = (
summary["last_close"] / summary["first_open"] - 1.0
)
summary["minute_intraday_range"] = summary["high"] / summary["low"] - 1.0
summary["minute_vwap"] = (
summary["amount_sum"] / summary["volume_sum"].where(summary["volume_sum"] > 0)
)
summary = summary.reset_index()
if daily is not None:
daily_keys = daily[["symbol_id", "date"]].copy()
daily_keys["date"] = pd.to_datetime(daily_keys["date"])
daily_keys = daily_keys.drop_duplicates(["symbol_id", "date"])
result = daily_keys.merge(summary, on=["symbol_id", "date"], how="left")
if "close" in daily.columns:
daily_close = daily[["symbol_id", "date", "close"]].copy()
daily_close["date"] = pd.to_datetime(daily_close["date"])
daily_close = daily_close.drop_duplicates(["symbol_id", "date"])
result = result.merge(
daily_close.rename(columns={"close": "daily_close"}),
on=["symbol_id", "date"],
how="left",
)
reference_close = result["daily_close"].fillna(result["last_close"])
else:
reference_close = result["last_close"]
else:
result = summary
reference_close = result["last_close"]
result["minute_vwap_deviation"] = (
result["minute_vwap"] / reference_close.replace(0.0, np.nan) - 1.0
)
return result[
[
"symbol_id",
"date",
"minute_bar_count",
"minute_intraday_return",
"minute_intraday_range",
"minute_vwap",
"minute_vwap_deviation",
]
]
+79
View File
@@ -0,0 +1,79 @@
"""Registry and factory for daily feature plugins."""
import importlib
import importlib.util
import inspect
from pathlib import Path
from typing import Optional, Type
from pipeline.features.base import BaseFeature
_REGISTRY: dict[str, Type[BaseFeature]] = {}
_builtins_loaded = False
def register_feature(cls: Type[BaseFeature]) -> Type[BaseFeature]:
"""Class decorator that registers a feature under ``BaseFeature.name``."""
if not (isinstance(cls, type) and issubclass(cls, BaseFeature)):
raise TypeError(f"{cls!r} is not a BaseFeature subclass")
key = getattr(cls, "name", "")
if not key:
raise ValueError(f"{cls.__name__} must set a non-empty class attribute `name`")
existing = _REGISTRY.get(key)
if existing is not None and existing is not cls:
raise ValueError(
f"Feature name '{key}' already registered by {existing.__name__}"
)
_REGISTRY[key] = cls
return cls
def available_features() -> list[str]:
"""Sorted names of all registered features (built-ins are loaded lazily)."""
_ensure_builtins()
return sorted(_REGISTRY)
def get_feature(name: str, **params) -> BaseFeature:
"""Instantiate a registered feature by name.
Only parameters accepted by the feature class's ``__init__`` are forwarded.
"""
_ensure_builtins()
if name not in _REGISTRY:
raise KeyError(f"Unknown feature '{name}'. Available: {sorted(_REGISTRY)}")
cls = _REGISTRY[name]
accepted = _accepted_params(cls)
kwargs = params if accepted is None else {k: v for k, v in params.items() if k in accepted}
return cls(**kwargs)
def load_feature_module(spec: str) -> None:
"""Import an external module so its ``@register_feature`` classes register."""
looks_like_file = spec.endswith(".py") or Path(spec).expanduser().exists()
if looks_like_file:
path = Path(spec).expanduser().resolve()
if not path.exists():
raise FileNotFoundError(f"Feature module not found: {path}")
module_spec = importlib.util.spec_from_file_location(path.stem, path)
if module_spec is None or module_spec.loader is None:
raise ImportError(f"Cannot load feature module from {path}")
module = importlib.util.module_from_spec(module_spec)
module_spec.loader.exec_module(module)
else:
importlib.import_module(spec)
def _accepted_params(cls: Type[BaseFeature]) -> Optional[set[str]]:
"""Param names ``cls.__init__`` accepts, or None if it takes ``**kwargs``."""
sig = inspect.signature(cls.__init__)
if any(p.kind is p.VAR_KEYWORD for p in sig.parameters.values()):
return None
return {name for name in sig.parameters if name != "self"}
def _ensure_builtins() -> None:
global _builtins_loaded
if not _builtins_loaded:
import pipeline.features.library # noqa: F401
_builtins_loaded = True