vizflow 0.4.4__py3-none-any.whl → 0.5.1__py3-none-any.whl

vizflow/__init__.py

@@ -5,7 +5,7 @@ Usage:
     import vizflow as vf
 """
 
-__version__ = "0.4.4"
+__version__ = "0.5.1"
 
 from .config import ColumnSchema, Config, get_config, set_config
 from .io import (
@@ -18,5 +18,5 @@ from .io import (
     scan_trades,
 )
 from .market import CN, CRYPTO, Market, Session
-from .ops import aggregate, bin, parse_time
+from .ops import aggregate, bin, forward_return, parse_time
 from .presets import JYAO_V20251114, PRESETS, YLIN_V20251204

vizflow/config.py

@@ -10,6 +10,21 @@ from typing import Any
 _global_config: Config | None = None
 
 
+def _validate_date(date: str) -> None:
+    """Validate date string format to prevent path traversal.
+
+    Args:
+        date: Date string to validate
+
+    Raises:
+        ValueError: If date is not exactly 8 digits (YYYYMMDD format)
+    """
+    if not (len(date) == 8 and date.isdigit()):
+        raise ValueError(
+            f"Invalid date format: {date!r}. Expected YYYYMMDD (8 digits)."
+        )
+
+
 @dataclass
 class ColumnSchema:
     """Schema for a column with type casting.
@@ -83,7 +98,11 @@ class Config:
     time_cutoff: int | None = None
 
     def __post_init__(self):
-        """Convert paths to Path objects if needed."""
+        """Convert string paths to Path objects.
+
+        Note: String values for path fields (alpha_dir, trade_dir, calendar_path,
+        replay_dir, aggregate_dir) are automatically converted to Path objects.
+        """
         if isinstance(self.alpha_dir, str):
             self.alpha_dir = Path(self.alpha_dir)
         if isinstance(self.trade_dir, str):
@@ -119,8 +138,9 @@ class Config:
             Full path to alpha file
 
         Raises:
-            ValueError: If alpha_dir is not set
+            ValueError: If alpha_dir is not set or date format is invalid
         """
+        _validate_date(date)
         if self.alpha_dir is None:
             raise ValueError("alpha_dir is not set in Config")
         return self.alpha_dir / self.alpha_pattern.format(date=date)
@@ -135,8 +155,9 @@ class Config:
             Full path to trade file
 
         Raises:
-            ValueError: If trade_dir is not set
+            ValueError: If trade_dir is not set or date format is invalid
         """
+        _validate_date(date)
         if self.trade_dir is None:
             raise ValueError("trade_dir is not set in Config")
         return self.trade_dir / self.trade_pattern.format(date=date)
@@ -152,8 +173,9 @@ class Config:
             Full path to replay output file
 
         Raises:
-            ValueError: If replay_dir is not set
+            ValueError: If replay_dir is not set or date format is invalid
         """
+        _validate_date(date)
         if self.replay_dir is None:
             raise ValueError("replay_dir is not set in Config")
         return self.replay_dir / f"{date}{suffix}"
@@ -169,8 +191,9 @@ class Config:
             Full path to aggregate output file
 
         Raises:
-            ValueError: If aggregate_dir is not set
+            ValueError: If aggregate_dir is not set or date format is invalid
         """
+        _validate_date(date)
         if self.aggregate_dir is None:
             raise ValueError("aggregate_dir is not set in Config")
         return self.aggregate_dir / f"{date}{suffix}"

vizflow/io.py

@@ -2,14 +2,24 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
-
 import polars as pl
 
-from .config import Config, get_config
+from .config import ColumnSchema, Config, get_config
+
 
-if TYPE_CHECKING:
-    pass
+def _apply_schema(df: pl.LazyFrame, schema: dict[str, ColumnSchema]) -> pl.LazyFrame:
+    """Apply type casting from schema definition.
+
+    Args:
+        df: LazyFrame to apply schema to
+        schema: Mapping from column names to ColumnSchema
+
+    Returns:
+        LazyFrame with type casts applied
+    """
+    for col_name, col_schema in schema.items():
+        df = df.with_columns(pl.col(col_name).cast(col_schema.cast_to))
+    return df
 
 
 def load_alpha(date: str, config: Config | None = None) -> pl.LazyFrame:
@@ -34,12 +44,7 @@ def load_alpha(date: str, config: Config | None = None) -> pl.LazyFrame:
     config = config or get_config()
     path = config.get_alpha_path(date)
     df = pl.scan_ipc(path)
-
-    # Apply schema evolution (type casting)
-    for col_name, schema in config.alpha_schema.items():
-        df = df.with_columns(pl.col(col_name).cast(schema.cast_to))
-
-    return df
+    return _apply_schema(df, config.alpha_schema)
 
 
 def load_trade(date: str, config: Config | None = None) -> pl.LazyFrame:
@@ -64,12 +69,7 @@ def load_trade(date: str, config: Config | None = None) -> pl.LazyFrame:
     config = config or get_config()
     path = config.get_trade_path(date)
     df = pl.scan_ipc(path)
-
-    # Apply schema evolution (type casting)
-    for col_name, schema in config.trade_schema.items():
-        df = df.with_columns(pl.col(col_name).cast(schema.cast_to))
-
-    return df
+    return _apply_schema(df, config.trade_schema)
 
 
 def load_calendar(config: Config | None = None) -> pl.DataFrame:
@@ -121,8 +121,10 @@ def _scan_file(path) -> pl.LazyFrame:
     elif suffix == "parquet":
         return pl.scan_parquet(path)
     else:
-        # Default to IPC
-        return pl.scan_ipc(path)
+        raise ValueError(
+            f"Unsupported file format: .{suffix}. "
+            "Supported: .feather, .ipc, .arrow, .csv, .meords, .parquet"
+        )
 
 
 def scan_trade(date: str, config: Config | None = None) -> pl.LazyFrame:
@@ -191,17 +193,13 @@ def scan_trades(config: Config | None = None) -> pl.LazyFrame:
 def _apply_trade_mapping(df: pl.LazyFrame, config: Config) -> pl.LazyFrame:
     """Apply column rename + schema evolution for trade data."""
     df = _apply_rename(df, config.trade_preset)
-    for col_name, schema in config.trade_schema.items():
-        df = df.with_columns(pl.col(col_name).cast(schema.cast_to))
-    return df
+    return _apply_schema(df, config.trade_schema)
 
 
 def _apply_alpha_mapping(df: pl.LazyFrame, config: Config) -> pl.LazyFrame:
     """Apply column rename + schema evolution for alpha data."""
     df = _apply_rename(df, config.alpha_preset)
-    for col_name, schema in config.alpha_schema.items():
-        df = df.with_columns(pl.col(col_name).cast(schema.cast_to))
-    return df
+    return _apply_schema(df, config.alpha_schema)
 
 
 def _apply_rename(df: pl.LazyFrame, preset: str | None) -> pl.LazyFrame:

vizflow/ops.py

@@ -135,3 +135,103 @@ def aggregate(
     """
     agg_exprs = [expr.alias(name) for name, expr in metrics.items()]
     return df.group_by(group_by).agg(agg_exprs)
+
+
+def _horizon_to_suffix(horizon_seconds: int) -> str:
+    """Convert horizon in seconds to column name suffix.
+
+    Rule: ≤60s → use seconds (60s), >60s → use minutes (3m, 30m)
+    """
+    if horizon_seconds <= 60:
+        return f"{horizon_seconds}s"
+    else:
+        minutes = horizon_seconds // 60
+        return f"{minutes}m"
+
+
+def forward_return(
+    df_trade: pl.LazyFrame,
+    df_alpha: pl.LazyFrame,
+    horizons: list[int],
+    trade_time_col: str = "elapsed_alpha_ts",
+    alpha_time_col: str = "elapsed_ticktime",
+    price_col: str = "mid",
+    symbol_col: str = "ukey",
+    tolerance_ms: int = 5000,
+) -> pl.LazyFrame:
+    """Merge alpha's future price to trade and calculate forward returns.
+
+    For each trade row:
+    1. Look up alpha price at trade_time + horizon
+    2. Add forward_{price_col}_{horizon} column (the future price)
+    3. Calculate y_{horizon} = (forward_price - current_price) / current_price
+
+    Output column names follow the convention:
+    - ≤60s → forward_mid_60s, y_60s
+    - >60s → forward_mid_3m, y_3m
+
+    Args:
+        df_trade: Trade LazyFrame with trade_time_col and price_col
+        df_alpha: Alpha LazyFrame with alpha_time_col and price_col
+        horizons: List of horizon in seconds, e.g., [60, 180, 1800]
+        trade_time_col: Time column in trade df (default: "elapsed_alpha_ts")
+        alpha_time_col: Time column in alpha df (default: "elapsed_ticktime")
+        price_col: Column name for price in both dfs (default: "mid")
+        symbol_col: Symbol column for grouping (default: "ukey")
+        tolerance_ms: Max time difference in ms for asof join (default: 5000)
+
+    Returns:
+        Trade LazyFrame with forward_* and y_* columns added
+
+    Example:
+        >>> df_trade = vf.parse_time(vf.scan_trade(date), "alpha_ts")
+        >>> df_alpha = vf.parse_time(vf.scan_alpha(date), "ticktime")
+        >>> df = vf.forward_return(df_trade, df_alpha, horizons=[60, 180, 1800])
+        >>> # Creates: forward_mid_60s, forward_mid_3m, forward_mid_30m
+        >>> #          y_60s, y_3m, y_30m
+    """
+    # Collect for asof join
+    trade = df_trade.collect()
+    alpha = df_alpha.collect()
+
+    # Prepare alpha lookup table: (symbol, time) -> price
+    alpha_lookup = alpha.select([
+        pl.col(symbol_col),
+        pl.col(alpha_time_col),
+        pl.col(price_col),
+    ]).sort([symbol_col, alpha_time_col])
+
+    for horizon in horizons:
+        suffix = _horizon_to_suffix(horizon)
+        horizon_ms = horizon * 1000
+        forward_col = f"forward_{price_col}_{suffix}"
+        return_col = f"y_{suffix}"
+
+        # Add target time column for this horizon
+        trade = trade.with_columns(
+            (pl.col(trade_time_col) + horizon_ms).alias("_forward_time")
+        )
+
+        # Sort by join columns (required for asof join)
+        trade = trade.sort([symbol_col, "_forward_time"])
+
+        # Asof join: find alpha price at forward_time
+        joined = trade.join_asof(
+            alpha_lookup.rename({alpha_time_col: "_alpha_time", price_col: "_forward_price"}),
+            left_on="_forward_time",
+            right_on="_alpha_time",
+            by=symbol_col,
+            strategy="nearest",
+            tolerance=tolerance_ms,
+        )
+
+        # Add forward price and calculate return (guard against zero price)
+        trade = joined.with_columns([
+            pl.col("_forward_price").alias(forward_col),
+            pl.when(pl.col(price_col) != 0)
+            .then((pl.col("_forward_price") - pl.col(price_col)) / pl.col(price_col))
+            .otherwise(pl.lit(None))
+            .alias(return_col),
+        ]).drop(["_forward_time", "_alpha_time", "_forward_price"])
+
+    return trade.lazy()

{vizflow-0.4.4.dist-info → vizflow-0.5.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: vizflow
-Version: 0.4.4
+Version: 0.5.1
 Requires-Python: >=3.9
 Requires-Dist: polars>=0.20.0
 Provides-Extra: dev

vizflow-0.5.1.dist-info/RECORD

@@ -0,0 +1,10 @@
+vizflow/__init__.py,sha256=EJ8qF4o2grf4aSochaasaaf0unyXV5yhrMs6rAhyp7k,496
+vizflow/config.py,sha256=JNW5-TshQ1v-Ft3-VV0JYJ5PdC3Yhgy4fW0AV0RWzkE,7322
+vizflow/io.py,sha256=ypiEuuPoHFKSt6VnhXcEI7u7dyVjKORunjqGpkFauXM,8877
+vizflow/market.py,sha256=MtNz_nnZxC66Aq-i2PXEwaFCTknijFWYZUUv6798k2s,2493
+vizflow/ops.py,sha256=4UwxOTPhvZ1_4PI3pcxbXfLAYsn1Ecj6nyBtBBr7KS8,7761
+vizflow/presets.py,sha256=h91NZoOH4YAx0bbsaNigECf9WOcWh1QZavguunWkaLE,2452
+vizflow/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+vizflow-0.5.1.dist-info/METADATA,sha256=DHzwPBvYuj7Rc4BawcXD2juS5iR5UD1FnAxt3cgvpo4,388
+vizflow-0.5.1.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+vizflow-0.5.1.dist-info/RECORD,,

vizflow-0.4.4.dist-info/RECORD

@@ -1,10 +0,0 @@
-vizflow/__init__.py,sha256=vYhGrFOQoONxBA6TJ08OQwWFd_UzltTsoZNQryNwnko,480
-vizflow/config.py,sha256=zSZnhdHzgXSqhDenHcHKm4CDGrMpKAdkNNRoUYYF1uc,6530
-vizflow/io.py,sha256=zmN0fFQOTmSBEBKangMExr0Q5mC2gajZM6GgdAyWkw4,8824
-vizflow/market.py,sha256=MtNz_nnZxC66Aq-i2PXEwaFCTknijFWYZUUv6798k2s,2493
-vizflow/ops.py,sha256=NL-Gtv-m_O1hv-0RUb9Wt43916HsQ5tYK_0e_uKR90w,4062
-vizflow/presets.py,sha256=h91NZoOH4YAx0bbsaNigECf9WOcWh1QZavguunWkaLE,2452
-vizflow/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-vizflow-0.4.4.dist-info/METADATA,sha256=NxMcTJ5fQKbB6GPak9dHARYTp9h0WwcDwxQAhrHcxRU,388
-vizflow-0.4.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-vizflow-0.4.4.dist-info/RECORD,,