From 58bef2c36d65e1daca45543b297dcc177eaf6e9d Mon Sep 17 00:00:00 2001 From: Thomas Schmelzer Date: Fri, 17 Jul 2026 23:07:53 +0400 Subject: [PATCH] refactor: split data.py god-module and decompose B(10) blocks (#871) data.py was the largest module in src/ (808 LOC) with the worst-ranking complexity blocks. This splits it and removes those hotspots: - Move the frame-construction/interpolation free functions into a new src/jquantstats/_utils/_construction.py (_to_polars, _apply_null_strategy, interpolate, _subtract_risk_free, _require_date_col, _align_returns_benchmark, plus a new _prices_to_returns). data.py re-imports them, so `jquantstats.data._subtract_risk_free` and `from jquantstats import interpolate` keep working. - Move the "derive a new Data" methods (resample, copy, head, tail, truncate and the _truncate_* helpers) into src/jquantstats/_data_reshape.py as _ReshapeMixin, which Data now inherits. The mixin carries empty __slots__ so the frozen slotted dataclass is unchanged; it rebuilds via a lazily-imported Data to avoid the accessor import cycle. - Decompose the B(10)/B(8) blocks: _apply_null_strategy -> B(5) (extracted _value_columns/_columns_with_nulls), Data.from_prices -> A(2) (extracted _prices_to_returns), _truncate_integer -> A (extracted _resolve_row_bound). Result: data.py drops 808 -> 453 LOC; no block in data.py exceeds B(6). make test (100% coverage), typecheck (ty + mypy strict) and docs-coverage (100%) all stay green. Closes #871 Co-Authored-By: Claude Opus 4.8 (1M context) --- src/jquantstats/_data_reshape.py | 219 +++++++++++++ src/jquantstats/_utils/_construction.py | 274 +++++++++++++++++ src/jquantstats/data.py | 391 +----------------------- 3 files changed, 508 insertions(+), 376 deletions(-) create mode 100644 src/jquantstats/_data_reshape.py create mode 100644 src/jquantstats/_utils/_construction.py diff --git a/src/jquantstats/_data_reshape.py b/src/jquantstats/_data_reshape.py new file mode 100644 index 00000000..ff5fd9d5 --- /dev/null +++ b/src/jquantstats/_data_reshape.py @@ -0,0 +1,219 @@ +"""Reshaping operations for `Data`: resampling, copying, slicing, truncation. + +`_ReshapeMixin` collects the methods that return a *new* `Data` derived from an +existing one. They are factored out of ``data.py`` to keep that module focused +on construction and accessors; the mixin only reads the three dataclass fields +(``returns``, ``index``, ``benchmark``) and rebuilds via `_rebuild`, which +imports `Data` lazily to avoid re-forming the ``data`` ↔ mixin import cycle. +""" + +from __future__ import annotations + +from datetime import date, datetime +from typing import TYPE_CHECKING + +import polars as pl + +from .exceptions import IntegerIndexBoundError + +if TYPE_CHECKING: + from .data import Data + + +class _ReshapeMixin: + """Mixin providing the `Data` operations that yield a new `Data`. + + The concrete class (`Data`) supplies the ``returns``, ``index`` and + ``benchmark`` dataclass fields; they are annotated here so the mixin's + methods type-check without importing `Data` at module load (which would + re-form an import cycle). No runtime attributes are created — the mixin + carries empty slots. + """ + + __slots__ = () + + # Provided by the concrete Data dataclass; declared for type-checkers only. + returns: pl.DataFrame + index: pl.DataFrame + benchmark: pl.DataFrame | None + + def _rebuild( + self, + *, + returns: pl.DataFrame, + index: pl.DataFrame, + benchmark: pl.DataFrame | None = None, + ) -> Data: + """Build a fresh `Data` from the given frames. + + `Data` is imported lazily so this module stays importable from + ``data.py`` without a cycle. + + Args: + returns: Returns frame for the new object. + index: Date/row index frame for the new object. + benchmark: Optional benchmark frame for the new object. + + Returns: + Data: A new `Data` built from the supplied frames. + """ + from .data import Data + + return Data(returns=returns, index=index, benchmark=benchmark) + + def resample(self, every: str = "1mo") -> Data: + """Resample returns and benchmark to a different frequency. + + Args: + every (str): Resampling frequency (e.g., ``'1mo'``, ``'1y'``). + Defaults to ``'1mo'``. + + Returns: + Data: Resampled data at the requested frequency. + + """ + + def resample_frame(dframe: pl.DataFrame) -> pl.DataFrame: + """Resample a single DataFrame to the target frequency using compound returns.""" + dframe = self.index.hstack(dframe) # Add the date column for resampling + + return dframe.group_by_dynamic( + index_column=self.index.columns[0], every=every, period=every, closed="right", label="right" + ).agg( + [ + ((pl.col(col) + 1.0).product() - 1.0).alias(col) + for col in dframe.columns + if col != self.index.columns[0] + ] + ) + + resampled_returns = resample_frame(self.returns) + resampled_benchmark = resample_frame(self.benchmark) if self.benchmark is not None else None + resampled_index = resampled_returns.select(self.index.columns[0]) + + return self._rebuild( + returns=resampled_returns.drop(self.index.columns[0]), + benchmark=resampled_benchmark.drop(self.index.columns[0]) if resampled_benchmark is not None else None, + index=resampled_index, + ) + + def copy(self) -> Data: + """Create a deep copy of the Data object. + + Returns: + Data: A new Data object with copies of the returns and benchmark. + + """ + benchmark = self.benchmark.clone() if self.benchmark is not None else None + return self._rebuild(returns=self.returns.clone(), benchmark=benchmark, index=self.index.clone()) + + def head(self, n: int = 5) -> Data: + """Return the first n rows of the combined returns and benchmark data. + + Args: + n (int, optional): Number of rows to return. Defaults to 5. + + Returns: + Data: A new Data object containing the first n rows of the combined data. + + """ + benchmark_head = self.benchmark.head(n) if self.benchmark is not None else None + return self._rebuild(returns=self.returns.head(n), benchmark=benchmark_head, index=self.index.head(n)) + + def tail(self, n: int = 5) -> Data: + """Return the last n rows of the combined returns and benchmark data. + + Args: + n (int, optional): Number of rows to return. Defaults to 5. + + Returns: + Data: A new Data object containing the last n rows of the combined data. + + """ + benchmark_tail = self.benchmark.tail(n) if self.benchmark is not None else None + return self._rebuild(returns=self.returns.tail(n), benchmark=benchmark_tail, index=self.index.tail(n)) + + def truncate( + self, + start: date | datetime | str | int | None = None, + end: date | datetime | str | int | None = None, + ) -> Data: + """Return a new Data object truncated to the inclusive [start, end] range. + + When the index is temporal (Date/Datetime), truncation is performed by + comparing the date column against ``start`` and ``end`` values. + + When the index is integer-based, row slicing is used instead, and + ``start`` and ``end`` must be non-negative integers. Passing + non-integer bounds to an integer-indexed Data raises `TypeError`. + + Args: + start: Optional lower bound (inclusive). A date/datetime value + when the index is temporal; a non-negative `int` row + index when the data has no temporal index. + end: Optional upper bound (inclusive). Same type rules as + ``start``. + + Returns: + Data: A new Data object filtered to the specified range. + + Raises: + TypeError: When the index is not temporal and a non-integer bound + is supplied. + + """ + date_column = self.index.columns[0] + + if self.index[date_column].dtype.is_temporal(): + new_index, new_returns, new_benchmark = self._truncate_temporal(date_column, start, end) + else: + new_index, new_returns, new_benchmark = self._truncate_integer(start, end) + + return self._rebuild(returns=new_returns, benchmark=new_benchmark, index=new_index) + + def _truncate_temporal( + self, + date_column: str, + start: date | datetime | str | int | None, + end: date | datetime | str | int | None, + ) -> tuple[pl.DataFrame, pl.DataFrame, pl.DataFrame | None]: + """Truncate a temporal index by comparing the date column to [start, end].""" + cond = pl.lit(True) + if start is not None: + cond = cond & (pl.col(date_column) >= pl.lit(start)) + if end is not None: + cond = cond & (pl.col(date_column) <= pl.lit(end)) + mask = self.index.select(cond.alias("mask"))["mask"] + new_benchmark = self.benchmark.filter(mask) if self.benchmark is not None else None + return self.index.filter(mask), self.returns.filter(mask), new_benchmark + + @staticmethod + def _resolve_row_bound(name: str, value: date | datetime | str | int | None, default: int) -> int: + """Validate and resolve an integer truncation bound to a row index. + + Args: + name: The bound's name (``"start"`` or ``"end"``) for the message. + value: The supplied bound; ``None`` and ``int`` are accepted. + default: The row index to use when *value* is ``None``. + + Returns: + int: *value* when it is an ``int``, otherwise *default*. + + Raises: + IntegerIndexBoundError: If *value* is neither ``None`` nor ``int``. + """ + if value is not None and not isinstance(value, int): + raise IntegerIndexBoundError(name, type(value).__name__) + return value if value is not None else default + + def _truncate_integer( + self, + start: date | datetime | str | int | None, + end: date | datetime | str | int | None, + ) -> tuple[pl.DataFrame, pl.DataFrame, pl.DataFrame | None]: + """Truncate an integer index by row slicing; bounds must be integers.""" + row_start = self._resolve_row_bound("start", start, 0) + row_end = self._resolve_row_bound("end", end, self.index.height - 1) + 1 + length = max(0, row_end - row_start) + new_benchmark = self.benchmark.slice(row_start, length) if self.benchmark is not None else None + return self.index.slice(row_start, length), self.returns.slice(row_start, length), new_benchmark diff --git a/src/jquantstats/_utils/_construction.py b/src/jquantstats/_utils/_construction.py new file mode 100644 index 00000000..c2e9ffd6 --- /dev/null +++ b/src/jquantstats/_utils/_construction.py @@ -0,0 +1,274 @@ +"""Frame-construction and interpolation helpers for building `Data` objects. + +These free functions carry the returns/prices ingestion logic that used to live +inline in ``data.py``: native-frame coercion, null handling, interior +interpolation, risk-free subtraction, date-column validation, and +returns/benchmark alignment. Keeping them here shrinks ``data.py`` and isolates +the construction concern from the `Data` container itself. +""" + +from __future__ import annotations + +import warnings +from typing import Literal + +import narwhals as nw +import polars as pl + +from .._types import NativeFrame +from ..exceptions import ( + BenchmarkAlignmentWarning, + MissingDateColumnError, + NullsInReturnsError, +) + +__all__ = [ + "interpolate", +] + + +def _to_polars(df: NativeFrame) -> pl.DataFrame: + """Convert any narwhals-compatible DataFrame to a polars DataFrame.""" + if isinstance(df, pl.DataFrame): + return df + return nw.from_native(df, eager_only=True).to_polars() + + +def _value_columns(dframe: pl.DataFrame, date_col: str) -> list[str]: + """Return every column of *dframe* except the date column.""" + return [c for c in dframe.columns if c != date_col] + + +def _columns_with_nulls(dframe: pl.DataFrame, value_cols: list[str]) -> list[str]: + """Return the subset of *value_cols* that contain at least one null.""" + null_counts = dframe.select(value_cols).null_count().row(0) + return [col for col, count in zip(value_cols, null_counts, strict=False) if count > 0] + + +def _apply_null_strategy( + dframe: pl.DataFrame, + date_col: str, + frame_name: str, + null_strategy: Literal["raise", "drop", "forward_fill"] | None, +) -> pl.DataFrame: + """Check for nulls in *dframe* and apply *null_strategy*. + + Args: + dframe (pl.DataFrame): DataFrame to inspect. The date column is + excluded from the null scan. + date_col (str): Name of the column to treat as the date index + (excluded from null check). + frame_name (str): Descriptive name used in the error message + (e.g. ``"returns"``). + null_strategy ({"raise", "drop", "forward_fill"} | None): How to + handle null values: + + - ``None`` — leave nulls as-is (nulls will propagate through + calculations). + - ``"raise"`` — raise `NullsInReturnsError` if any null is found. + - ``"drop"`` — drop every row that contains at least one null. + - ``"forward_fill"`` — fill each null with the most recent + non-null value in the same column. + + Returns: + pl.DataFrame: The original DataFrame (``None`` / ``"raise"``), a + filtered DataFrame (``"drop"``), or a filled DataFrame + (``"forward_fill"``). + + Raises: + NullsInReturnsError: When *null_strategy* is ``"raise"`` and nulls + are present. + + """ + if null_strategy is None: + return dframe + + value_cols = _value_columns(dframe, date_col) + cols_with_nulls = _columns_with_nulls(dframe, value_cols) + + if not cols_with_nulls: + return dframe + + if null_strategy == "raise": + raise NullsInReturnsError(frame_name, cols_with_nulls) + if null_strategy == "drop": + return dframe.drop_nulls(subset=value_cols) + # forward_fill + return dframe.with_columns(pl.col(value_cols).forward_fill()) + + +def interpolate(df: pl.DataFrame) -> pl.DataFrame: + """Forward-fill numeric columns only between first and last non-null values. + + For each numeric column, forward-fill is applied strictly within the span + bounded by its first and last non-null samples. Values outside this span + are left as-is (including leading/trailing nulls). Non-numeric columns are + returned unchanged. + + Args: + df: Input frame possibly containing nulls. + + Returns: + pl.DataFrame: Frame where numeric columns have been interior-forward- + filled; schema and dtypes of the original columns are preserved. + + Examples: + ```python + import polars as pl + from jquantstats import interpolate + + df = pl.DataFrame({"a": [None, 1.0, None, 3.0, None], "b": ["x", "y", "z", "w", "v"]}) + result = interpolate(df) + # a: [None, 1.0, 1.0, 3.0, None] (leading/trailing nulls untouched) + # b: ["x", "y", "z", "w", "v"] (non-numeric unchanged) + ``` + + """ + # Choose a temp column name guaranteed not to collide with any user column. + tmp_col = "__row_idx__" + while tmp_col in df.columns: + tmp_col = f"_{tmp_col}_" + + out = [] + + for col in df.columns: + s = df[col] + if s.dtype.is_numeric(): + non_null_mask = s.is_not_null() + if non_null_mask.any(): + _fwd = non_null_mask.arg_max() + _rev = non_null_mask.reverse().arg_max() + if _fwd is None or _rev is None: # pragma: no cover + out.append(pl.col(col)) + continue + first_valid_idx = _fwd + last_valid_idx = len(s) - 1 - _rev + else: + out.append(pl.col(col)) + continue + + mask = (pl.col(tmp_col) >= pl.lit(first_valid_idx)) & (pl.col(tmp_col) <= pl.lit(last_valid_idx)) + filled_col = pl.when(mask).then(pl.col(col).fill_null(strategy="forward")).otherwise(pl.col(col)).alias(col) + out.append(filled_col) + else: + out.append(pl.col(col)) + + return df.with_columns(pl.int_range(0, df.height).alias(tmp_col)).select(out) + + +def _subtract_risk_free(dframe: pl.DataFrame, rf: float | pl.DataFrame, date_col: str) -> pl.DataFrame: + """Subtract the risk-free rate from all numeric columns in the DataFrame. + + Args: + dframe (pl.DataFrame): DataFrame containing returns data with a date + column and one or more numeric columns representing asset returns. + rf (float | pl.DataFrame): Risk-free rate to subtract from returns. + + - If float: A constant risk-free rate applied to all dates. + - If pl.DataFrame: A DataFrame with a date column and a second + column containing time-varying risk-free rates. + + date_col (str): Name of the date column in both DataFrames for + joining when rf is a DataFrame. + + Returns: + pl.DataFrame: DataFrame with the risk-free rate subtracted from all + numeric columns, preserving the original column names. + + """ + if isinstance(rf, float): + rf_dframe = dframe.select([pl.col(date_col), pl.lit(rf).alias("rf")]) + else: + if not isinstance(rf, pl.DataFrame): + raise TypeError("rf must be a float or DataFrame") # noqa: TRY003 + if rf.columns[1] != "rf": + warnings.warn( + f"Risk-free rate column '{rf.columns[1]}' has been renamed to 'rf' for internal alignment.", + stacklevel=3, + ) + rf_dframe = rf.rename({rf.columns[1]: "rf"}) if rf.columns[1] != "rf" else rf + + dframe = dframe.join(rf_dframe, on=date_col, how="inner") + return dframe.select( + [pl.col(date_col)] + + [(pl.col(col) - pl.col("rf")).alias(col) for col in dframe.columns if col not in {date_col, "rf"}] + ) + + +def _require_date_col(frames: list[tuple[str, pl.DataFrame | None]], date_col: str) -> None: + """Verify *date_col* is present in every supplied (non-None) frame. + + Args: + frames: ``(name, frame)`` pairs; ``None`` frames are skipped. + date_col: The required date column name. + + Raises: + MissingDateColumnError: If any frame lacks *date_col*, naming that frame. + """ + for frame_name, frame in frames: + if frame is not None and date_col not in frame.columns: + raise MissingDateColumnError(frame_name, column=date_col, available=list(frame.columns)) + + +def _align_returns_benchmark( + returns_pl: pl.DataFrame, benchmark_pl: pl.DataFrame, date_col: str +) -> tuple[pl.DataFrame, pl.DataFrame]: + """Inner-join returns and benchmark on their common dates. + + Args: + returns_pl: Returns frame with a *date_col* column. + benchmark_pl: Benchmark frame with a *date_col* column. + date_col: The shared date column name. + + Returns: + The two frames filtered to their overlapping dates. + + Raises: + ValueError: If the frames share no dates. + + Warns: + BenchmarkAlignmentWarning: If aligning drops rows from either frame. + """ + joined_dates = returns_pl.join(benchmark_pl, on=date_col, how="inner").select(date_col) + if joined_dates.is_empty(): + raise ValueError("No overlapping dates between returns and benchmark.") # noqa: TRY003 + dropped_returns = returns_pl.height - joined_dates.height + dropped_benchmark = benchmark_pl.height - joined_dates.height + if dropped_returns > 0 or dropped_benchmark > 0: + warnings.warn( + f"Aligning returns and benchmark on common dates dropped " + f"{dropped_returns} of {returns_pl.height} returns row(s) and " + f"{dropped_benchmark} of {benchmark_pl.height} benchmark row(s); " + f"{joined_dates.height} row(s) remain. Pass a benchmark covering " + f"the same dates as the returns to avoid this.", + BenchmarkAlignmentWarning, + stacklevel=2, + ) + returns_pl = returns_pl.join(joined_dates, on=date_col, how="inner") + benchmark_pl = benchmark_pl.join(joined_dates, on=date_col, how="inner") + return returns_pl, benchmark_pl + + +def _prices_to_returns(frame: pl.DataFrame, date_col: str, frame_name: str) -> pl.DataFrame: + """Convert a price-level frame to a returns frame via percentage change. + + The first row is dropped because no prior price is available to compute a + return for it. + + Args: + frame: Price-level frame with a *date_col* column and asset columns. + date_col: Name of the date column (passed through unchanged). + frame_name: Descriptive name used in the error message when *date_col* + is missing (e.g. ``"prices"`` or ``"benchmark"``). + + Returns: + pl.DataFrame: Returns frame with the same columns as *frame*, one row + shorter. + + Raises: + MissingDateColumnError: If *date_col* is not a column of *frame*. + """ + if date_col not in frame.columns: + raise MissingDateColumnError(frame_name, column=date_col, available=list(frame.columns)) + asset_cols = _value_columns(frame, date_col) + return frame.with_columns([pl.col(c).pct_change().alias(c) for c in asset_cols]).slice(1) diff --git a/src/jquantstats/data.py b/src/jquantstats/data.py index 76bc5e99..a85059e3 100644 --- a/src/jquantstats/data.py +++ b/src/jquantstats/data.py @@ -3,20 +3,24 @@ from __future__ import annotations import dataclasses -import warnings from collections.abc import Iterator -from datetime import date, datetime, timedelta +from datetime import timedelta from typing import TYPE_CHECKING, Literal, cast -import narwhals as nw import polars as pl +from ._data_reshape import _ReshapeMixin from ._types import NativeFrame, NativeFrameOrScalar -from .exceptions import ( - BenchmarkAlignmentWarning, - IntegerIndexBoundError, - MissingDateColumnError, - NullsInReturnsError, +from ._utils._construction import ( + _align_returns_benchmark, + _apply_null_strategy, + _prices_to_returns, + _require_date_col, + _subtract_risk_free, + _to_polars, +) +from ._utils._construction import ( + interpolate as interpolate, # re-exported for `from jquantstats import interpolate` ) if TYPE_CHECKING: @@ -26,220 +30,8 @@ from ._utils import DataUtils -def _to_polars(df: NativeFrame) -> pl.DataFrame: - """Convert any narwhals-compatible DataFrame to a polars DataFrame.""" - if isinstance(df, pl.DataFrame): - return df - return nw.from_native(df, eager_only=True).to_polars() - - -def _apply_null_strategy( - dframe: pl.DataFrame, - date_col: str, - frame_name: str, - null_strategy: Literal["raise", "drop", "forward_fill"] | None, -) -> pl.DataFrame: - """Check for nulls in *dframe* and apply *null_strategy*. - - Args: - dframe (pl.DataFrame): DataFrame to inspect. The date column is - excluded from the null scan. - date_col (str): Name of the column to treat as the date index - (excluded from null check). - frame_name (str): Descriptive name used in the error message - (e.g. ``"returns"``). - null_strategy ({"raise", "drop", "forward_fill"} | None): How to - handle null values: - - - ``None`` — leave nulls as-is (nulls will propagate through - calculations). - - ``"raise"`` — raise `NullsInReturnsError` if any null is found. - - ``"drop"`` — drop every row that contains at least one null. - - ``"forward_fill"`` — fill each null with the most recent - non-null value in the same column. - - Returns: - pl.DataFrame: The original DataFrame (``None`` / ``"raise"``), a - filtered DataFrame (``"drop"``), or a filled DataFrame - (``"forward_fill"``). - - Raises: - NullsInReturnsError: When *null_strategy* is ``"raise"`` and nulls - are present. - - """ - if null_strategy is None: - return dframe - - value_cols = [c for c in dframe.columns if c != date_col] - null_counts = dframe.select(value_cols).null_count().row(0) - cols_with_nulls = [col for col, count in zip(value_cols, null_counts, strict=False) if count > 0] - - if not cols_with_nulls: - return dframe - - if null_strategy == "raise": - raise NullsInReturnsError(frame_name, cols_with_nulls) - if null_strategy == "drop": - return dframe.drop_nulls(subset=value_cols) - # forward_fill - return dframe.with_columns([pl.col(c).forward_fill() for c in value_cols]) - - -def interpolate(df: pl.DataFrame) -> pl.DataFrame: - """Forward-fill numeric columns only between first and last non-null values. - - For each numeric column, forward-fill is applied strictly within the span - bounded by its first and last non-null samples. Values outside this span - are left as-is (including leading/trailing nulls). Non-numeric columns are - returned unchanged. - - Args: - df: Input frame possibly containing nulls. - - Returns: - pl.DataFrame: Frame where numeric columns have been interior-forward- - filled; schema and dtypes of the original columns are preserved. - - Examples: - ```python - import polars as pl - from jquantstats import interpolate - - df = pl.DataFrame({"a": [None, 1.0, None, 3.0, None], "b": ["x", "y", "z", "w", "v"]}) - result = interpolate(df) - # a: [None, 1.0, 1.0, 3.0, None] (leading/trailing nulls untouched) - # b: ["x", "y", "z", "w", "v"] (non-numeric unchanged) - ``` - - """ - # Choose a temp column name guaranteed not to collide with any user column. - tmp_col = "__row_idx__" - while tmp_col in df.columns: - tmp_col = f"_{tmp_col}_" - - out = [] - - for col in df.columns: - s = df[col] - if s.dtype.is_numeric(): - non_null_mask = s.is_not_null() - if non_null_mask.any(): - _fwd = non_null_mask.arg_max() - _rev = non_null_mask.reverse().arg_max() - if _fwd is None or _rev is None: # pragma: no cover - out.append(pl.col(col)) - continue - first_valid_idx = _fwd - last_valid_idx = len(s) - 1 - _rev - else: - out.append(pl.col(col)) - continue - - mask = (pl.col(tmp_col) >= pl.lit(first_valid_idx)) & (pl.col(tmp_col) <= pl.lit(last_valid_idx)) - filled_col = pl.when(mask).then(pl.col(col).fill_null(strategy="forward")).otherwise(pl.col(col)).alias(col) - out.append(filled_col) - else: - out.append(pl.col(col)) - - return df.with_columns(pl.int_range(0, df.height).alias(tmp_col)).select(out) - - -def _subtract_risk_free(dframe: pl.DataFrame, rf: float | pl.DataFrame, date_col: str) -> pl.DataFrame: - """Subtract the risk-free rate from all numeric columns in the DataFrame. - - Args: - dframe (pl.DataFrame): DataFrame containing returns data with a date - column and one or more numeric columns representing asset returns. - rf (float | pl.DataFrame): Risk-free rate to subtract from returns. - - - If float: A constant risk-free rate applied to all dates. - - If pl.DataFrame: A DataFrame with a date column and a second - column containing time-varying risk-free rates. - - date_col (str): Name of the date column in both DataFrames for - joining when rf is a DataFrame. - - Returns: - pl.DataFrame: DataFrame with the risk-free rate subtracted from all - numeric columns, preserving the original column names. - - """ - if isinstance(rf, float): - rf_dframe = dframe.select([pl.col(date_col), pl.lit(rf).alias("rf")]) - else: - if not isinstance(rf, pl.DataFrame): - raise TypeError("rf must be a float or DataFrame") # noqa: TRY003 - if rf.columns[1] != "rf": - warnings.warn( - f"Risk-free rate column '{rf.columns[1]}' has been renamed to 'rf' for internal alignment.", - stacklevel=3, - ) - rf_dframe = rf.rename({rf.columns[1]: "rf"}) if rf.columns[1] != "rf" else rf - - dframe = dframe.join(rf_dframe, on=date_col, how="inner") - return dframe.select( - [pl.col(date_col)] - + [(pl.col(col) - pl.col("rf")).alias(col) for col in dframe.columns if col not in {date_col, "rf"}] - ) - - -def _require_date_col(frames: list[tuple[str, pl.DataFrame | None]], date_col: str) -> None: - """Verify *date_col* is present in every supplied (non-None) frame. - - Args: - frames: ``(name, frame)`` pairs; ``None`` frames are skipped. - date_col: The required date column name. - - Raises: - MissingDateColumnError: If any frame lacks *date_col*, naming that frame. - """ - for frame_name, frame in frames: - if frame is not None and date_col not in frame.columns: - raise MissingDateColumnError(frame_name, column=date_col, available=list(frame.columns)) - - -def _align_returns_benchmark( - returns_pl: pl.DataFrame, benchmark_pl: pl.DataFrame, date_col: str -) -> tuple[pl.DataFrame, pl.DataFrame]: - """Inner-join returns and benchmark on their common dates. - - Args: - returns_pl: Returns frame with a *date_col* column. - benchmark_pl: Benchmark frame with a *date_col* column. - date_col: The shared date column name. - - Returns: - The two frames filtered to their overlapping dates. - - Raises: - ValueError: If the frames share no dates. - - Warns: - BenchmarkAlignmentWarning: If aligning drops rows from either frame. - """ - joined_dates = returns_pl.join(benchmark_pl, on=date_col, how="inner").select(date_col) - if joined_dates.is_empty(): - raise ValueError("No overlapping dates between returns and benchmark.") # noqa: TRY003 - dropped_returns = returns_pl.height - joined_dates.height - dropped_benchmark = benchmark_pl.height - joined_dates.height - if dropped_returns > 0 or dropped_benchmark > 0: - warnings.warn( - f"Aligning returns and benchmark on common dates dropped " - f"{dropped_returns} of {returns_pl.height} returns row(s) and " - f"{dropped_benchmark} of {benchmark_pl.height} benchmark row(s); " - f"{joined_dates.height} row(s) remain. Pass a benchmark covering " - f"the same dates as the returns to avoid this.", - BenchmarkAlignmentWarning, - stacklevel=2, - ) - returns_pl = returns_pl.join(joined_dates, on=date_col, how="inner") - benchmark_pl = benchmark_pl.join(joined_dates, on=date_col, how="inner") - return returns_pl, benchmark_pl - - @dataclasses.dataclass(frozen=True, slots=True) -class Data: +class Data(_ReshapeMixin): """A container for financial returns data and an optional benchmark. Provides methods for analyzing and manipulating financial returns data, @@ -477,21 +269,11 @@ def from_prices( ``` """ - prices_pl = _to_polars(prices) - if date_col not in prices_pl.columns: - raise MissingDateColumnError("prices", column=date_col, available=list(prices_pl.columns)) - asset_cols = [c for c in prices_pl.columns if c != date_col] - returns_pl = prices_pl.with_columns([pl.col(c).pct_change().alias(c) for c in asset_cols]).slice(1) + returns_pl = _prices_to_returns(_to_polars(prices), date_col, "prices") benchmark_returns: NativeFrame | None = None if benchmark is not None: - benchmark_pl = _to_polars(benchmark) - if date_col not in benchmark_pl.columns: - raise MissingDateColumnError("benchmark", column=date_col, available=list(benchmark_pl.columns)) - bench_cols = [c for c in benchmark_pl.columns if c != date_col] - benchmark_returns = benchmark_pl.with_columns([pl.col(c).pct_change().alias(c) for c in bench_cols]).slice( - 1 - ) + benchmark_returns = _prices_to_returns(_to_polars(benchmark), date_col, "benchmark") return cls.from_returns( returns=returns_pl, @@ -607,42 +389,6 @@ def all(self) -> pl.DataFrame: else: return pl.concat([self.index, self.returns, self.benchmark], how="horizontal_extend") - def resample(self, every: str = "1mo") -> Data: - """Resample returns and benchmark to a different frequency. - - Args: - every (str): Resampling frequency (e.g., ``'1mo'``, ``'1y'``). - Defaults to ``'1mo'``. - - Returns: - Data: Resampled data at the requested frequency. - - """ - - def resample_frame(dframe: pl.DataFrame) -> pl.DataFrame: - """Resample a single DataFrame to the target frequency using compound returns.""" - dframe = self.index.hstack(dframe) # Add the date column for resampling - - return dframe.group_by_dynamic( - index_column=self.index.columns[0], every=every, period=every, closed="right", label="right" - ).agg( - [ - ((pl.col(col) + 1.0).product() - 1.0).alias(col) - for col in dframe.columns - if col != self.index.columns[0] - ] - ) - - resampled_returns = resample_frame(self.returns) - resampled_benchmark = resample_frame(self.benchmark) if self.benchmark is not None else None - resampled_index = resampled_returns.select(self.index.columns[0]) - - return Data( - returns=resampled_returns.drop(self.index.columns[0]), - benchmark=resampled_benchmark.drop(self.index.columns[0]) if resampled_benchmark is not None else None, - index=resampled_index, - ) - def describe(self) -> pl.DataFrame: """Return a tidy summary of shape, date range and asset names. @@ -665,113 +411,6 @@ def describe(self) -> pl.DataFrame: } ) - def copy(self) -> Data: - """Create a deep copy of the Data object. - - Returns: - Data: A new Data object with copies of the returns and benchmark. - - """ - if self.benchmark is not None: - return Data(returns=self.returns.clone(), benchmark=self.benchmark.clone(), index=self.index.clone()) - return Data(returns=self.returns.clone(), index=self.index.clone()) - - def head(self, n: int = 5) -> Data: - """Return the first n rows of the combined returns and benchmark data. - - Args: - n (int, optional): Number of rows to return. Defaults to 5. - - Returns: - Data: A new Data object containing the first n rows of the combined data. - - """ - benchmark_head = self.benchmark.head(n) if self.benchmark is not None else None - return Data(returns=self.returns.head(n), benchmark=benchmark_head, index=self.index.head(n)) - - def tail(self, n: int = 5) -> Data: - """Return the last n rows of the combined returns and benchmark data. - - Args: - n (int, optional): Number of rows to return. Defaults to 5. - - Returns: - Data: A new Data object containing the last n rows of the combined data. - - """ - benchmark_tail = self.benchmark.tail(n) if self.benchmark is not None else None - return Data(returns=self.returns.tail(n), benchmark=benchmark_tail, index=self.index.tail(n)) - - def truncate( - self, - start: date | datetime | str | int | None = None, - end: date | datetime | str | int | None = None, - ) -> Data: - """Return a new Data object truncated to the inclusive [start, end] range. - - When the index is temporal (Date/Datetime), truncation is performed by - comparing the date column against ``start`` and ``end`` values. - - When the index is integer-based, row slicing is used instead, and - ``start`` and ``end`` must be non-negative integers. Passing - non-integer bounds to an integer-indexed Data raises `TypeError`. - - Args: - start: Optional lower bound (inclusive). A date/datetime value - when the index is temporal; a non-negative `int` row - index when the data has no temporal index. - end: Optional upper bound (inclusive). Same type rules as - ``start``. - - Returns: - Data: A new Data object filtered to the specified range. - - Raises: - TypeError: When the index is not temporal and a non-integer bound - is supplied. - - """ - date_column = self.index.columns[0] - - if self.index[date_column].dtype.is_temporal(): - new_index, new_returns, new_benchmark = self._truncate_temporal(date_column, start, end) - else: - new_index, new_returns, new_benchmark = self._truncate_integer(start, end) - - return Data(returns=new_returns, benchmark=new_benchmark, index=new_index) - - def _truncate_temporal( - self, - date_column: str, - start: date | datetime | str | int | None, - end: date | datetime | str | int | None, - ) -> tuple[pl.DataFrame, pl.DataFrame, pl.DataFrame | None]: - """Truncate a temporal index by comparing the date column to [start, end].""" - cond = pl.lit(True) - if start is not None: - cond = cond & (pl.col(date_column) >= pl.lit(start)) - if end is not None: - cond = cond & (pl.col(date_column) <= pl.lit(end)) - mask = self.index.select(cond.alias("mask"))["mask"] - new_benchmark = self.benchmark.filter(mask) if self.benchmark is not None else None - return self.index.filter(mask), self.returns.filter(mask), new_benchmark - - def _truncate_integer( - self, - start: date | datetime | str | int | None, - end: date | datetime | str | int | None, - ) -> tuple[pl.DataFrame, pl.DataFrame, pl.DataFrame | None]: - """Truncate an integer index by row slicing; bounds must be integers.""" - if start is not None and not isinstance(start, int): - raise IntegerIndexBoundError("start", type(start).__name__) - if end is not None and not isinstance(end, int): - raise IntegerIndexBoundError("end", type(end).__name__) - row_start = start if start is not None else 0 - row_end = end + 1 if end is not None else self.index.height - length = max(0, row_end - row_start) - new_benchmark = self.benchmark.slice(row_start, length) if self.benchmark is not None else None - return self.index.slice(row_start, length), self.returns.slice(row_start, length), new_benchmark - @property def _periods_per_year(self) -> float: """Estimate the number of periods per year based on average frequency in the index.