Coverage for src/jquantstats/_utils/_construction.py: 100%
80 statements
« prev ^ index » next coverage.py v7.15.3, created at 2026-08-06 04:52 +0000
1"""Frame-construction and interpolation helpers for building `Data` objects.
3These free functions carry the returns/prices ingestion logic that used to live
4inline in ``data.py``: native-frame coercion, null handling, interior
5interpolation, risk-free subtraction, date-column validation, and
6returns/benchmark alignment. Keeping them here shrinks ``data.py`` and isolates
7the construction concern from the `Data` container itself.
8"""
10from __future__ import annotations
12import warnings
13from typing import Literal
15import narwhals as nw
16import polars as pl
18from .._types import NativeFrame
19from ..exceptions import (
20 BenchmarkAlignmentWarning,
21 MissingDateColumnError,
22 NullsInReturnsError,
23)
25__all__ = [
26 "interpolate",
27]
30def _to_polars(df: NativeFrame) -> pl.DataFrame:
31 """Convert any narwhals-compatible DataFrame to a polars DataFrame."""
32 if isinstance(df, pl.DataFrame):
33 return df
34 return nw.from_native(df, eager_only=True).to_polars()
37def _value_columns(dframe: pl.DataFrame, date_col: str) -> list[str]:
38 """Return every column of *dframe* except the date column."""
39 return [c for c in dframe.columns if c != date_col]
42def _columns_with_nulls(dframe: pl.DataFrame, value_cols: list[str]) -> list[str]:
43 """Return the subset of *value_cols* that contain at least one null."""
44 null_counts = dframe.select(value_cols).null_count().row(0)
45 return [col for col, count in zip(value_cols, null_counts, strict=False) if count > 0]
48def _apply_null_strategy(
49 dframe: pl.DataFrame,
50 date_col: str,
51 frame_name: str,
52 null_strategy: Literal["raise", "drop", "forward_fill"] | None,
53) -> pl.DataFrame:
54 """Check for nulls in *dframe* and apply *null_strategy*.
56 Args:
57 dframe (pl.DataFrame): DataFrame to inspect. The date column is
58 excluded from the null scan.
59 date_col (str): Name of the column to treat as the date index
60 (excluded from null check).
61 frame_name (str): Descriptive name used in the error message
62 (e.g. ``"returns"``).
63 null_strategy ({"raise", "drop", "forward_fill"} | None): How to
64 handle null values:
66 - ``None`` — leave nulls as-is (nulls will propagate through
67 calculations).
68 - ``"raise"`` — raise `NullsInReturnsError` if any null is found.
69 - ``"drop"`` — drop every row that contains at least one null.
70 - ``"forward_fill"`` — fill each null with the most recent
71 non-null value in the same column.
73 Returns:
74 pl.DataFrame: The original DataFrame (``None`` / ``"raise"``), a
75 filtered DataFrame (``"drop"``), or a filled DataFrame
76 (``"forward_fill"``).
78 Raises:
79 NullsInReturnsError: When *null_strategy* is ``"raise"`` and nulls
80 are present.
82 """
83 if null_strategy is None:
84 return dframe
86 value_cols = _value_columns(dframe, date_col)
87 cols_with_nulls = _columns_with_nulls(dframe, value_cols)
89 if not cols_with_nulls:
90 return dframe
92 if null_strategy == "raise":
93 raise NullsInReturnsError(frame_name, cols_with_nulls)
94 if null_strategy == "drop":
95 return dframe.drop_nulls(subset=value_cols)
96 # forward_fill
97 return dframe.with_columns(pl.col(value_cols).forward_fill())
100def interpolate(df: pl.DataFrame) -> pl.DataFrame:
101 """Forward-fill numeric columns only between first and last non-null values.
103 For each numeric column, forward-fill is applied strictly within the span
104 bounded by its first and last non-null samples. Values outside this span
105 are left as-is (including leading/trailing nulls). Non-numeric columns are
106 returned unchanged.
108 Args:
109 df: Input frame possibly containing nulls.
111 Returns:
112 pl.DataFrame: Frame where numeric columns have been interior-forward-
113 filled; schema and dtypes of the original columns are preserved.
115 Examples:
116 ```python
117 import polars as pl
118 from jquantstats import interpolate
120 df = pl.DataFrame({"a": [None, 1.0, None, 3.0, None], "b": ["x", "y", "z", "w", "v"]})
121 result = interpolate(df)
122 # a: [None, 1.0, 1.0, 3.0, None] (leading/trailing nulls untouched)
123 # b: ["x", "y", "z", "w", "v"] (non-numeric unchanged)
124 ```
126 """
127 # Choose a temp column name guaranteed not to collide with any user column.
128 tmp_col = "__row_idx__"
129 while tmp_col in df.columns:
130 tmp_col = f"_{tmp_col}_"
132 out = []
134 for col in df.columns:
135 s = df[col]
136 if s.dtype.is_numeric():
137 non_null_mask = s.is_not_null()
138 if non_null_mask.any():
139 _fwd = non_null_mask.arg_max()
140 _rev = non_null_mask.reverse().arg_max()
141 if _fwd is None or _rev is None: # pragma: no cover
142 out.append(pl.col(col))
143 continue
144 first_valid_idx = _fwd
145 last_valid_idx = len(s) - 1 - _rev
146 else:
147 out.append(pl.col(col))
148 continue
150 mask = (pl.col(tmp_col) >= pl.lit(first_valid_idx)) & (pl.col(tmp_col) <= pl.lit(last_valid_idx))
151 filled_col = pl.when(mask).then(pl.col(col).fill_null(strategy="forward")).otherwise(pl.col(col)).alias(col)
152 out.append(filled_col)
153 else:
154 out.append(pl.col(col))
156 return df.with_columns(pl.int_range(0, df.height).alias(tmp_col)).select(out)
159def _subtract_risk_free(dframe: pl.DataFrame, rf: float | pl.DataFrame, date_col: str) -> pl.DataFrame:
160 """Subtract the risk-free rate from all numeric columns in the DataFrame.
162 Args:
163 dframe (pl.DataFrame): DataFrame containing returns data with a date
164 column and one or more numeric columns representing asset returns.
165 rf (float | pl.DataFrame): Risk-free rate to subtract from returns.
167 - If float: A constant risk-free rate applied to all dates.
168 - If pl.DataFrame: A DataFrame with a date column and a second
169 column containing time-varying risk-free rates.
171 date_col (str): Name of the date column in both DataFrames for
172 joining when rf is a DataFrame.
174 Returns:
175 pl.DataFrame: DataFrame with the risk-free rate subtracted from all
176 numeric columns, preserving the original column names.
178 """
179 if isinstance(rf, float):
180 rf_dframe = dframe.select([pl.col(date_col), pl.lit(rf).alias("rf")])
181 else:
182 if not isinstance(rf, pl.DataFrame):
183 raise TypeError("rf must be a float or DataFrame") # noqa: TRY003
184 if rf.columns[1] != "rf":
185 warnings.warn(
186 f"Risk-free rate column '{rf.columns[1]}' has been renamed to 'rf' for internal alignment.",
187 stacklevel=3,
188 )
189 rf_dframe = rf.rename({rf.columns[1]: "rf"}) if rf.columns[1] != "rf" else rf
191 dframe = dframe.join(rf_dframe, on=date_col, how="inner")
192 return dframe.select(
193 [pl.col(date_col)]
194 + [(pl.col(col) - pl.col("rf")).alias(col) for col in dframe.columns if col not in {date_col, "rf"}]
195 )
198def _require_date_col(frames: list[tuple[str, pl.DataFrame | None]], date_col: str) -> None:
199 """Verify *date_col* is present in every supplied (non-None) frame.
201 Args:
202 frames: ``(name, frame)`` pairs; ``None`` frames are skipped.
203 date_col: The required date column name.
205 Raises:
206 MissingDateColumnError: If any frame lacks *date_col*, naming that frame.
207 """
208 for frame_name, frame in frames:
209 if frame is not None and date_col not in frame.columns:
210 raise MissingDateColumnError(frame_name, column=date_col, available=list(frame.columns))
213def _align_returns_benchmark(
214 returns_pl: pl.DataFrame, benchmark_pl: pl.DataFrame, date_col: str
215) -> tuple[pl.DataFrame, pl.DataFrame]:
216 """Inner-join returns and benchmark on their common dates.
218 Args:
219 returns_pl: Returns frame with a *date_col* column.
220 benchmark_pl: Benchmark frame with a *date_col* column.
221 date_col: The shared date column name.
223 Returns:
224 The two frames filtered to their overlapping dates.
226 Raises:
227 ValueError: If the frames share no dates.
229 Warns:
230 BenchmarkAlignmentWarning: If aligning drops rows from either frame.
231 """
232 joined_dates = returns_pl.join(benchmark_pl, on=date_col, how="inner").select(date_col)
233 if joined_dates.is_empty():
234 raise ValueError("No overlapping dates between returns and benchmark.") # noqa: TRY003
235 dropped_returns = returns_pl.height - joined_dates.height
236 dropped_benchmark = benchmark_pl.height - joined_dates.height
237 if dropped_returns > 0 or dropped_benchmark > 0:
238 warnings.warn(
239 f"Aligning returns and benchmark on common dates dropped "
240 f"{dropped_returns} of {returns_pl.height} returns row(s) and "
241 f"{dropped_benchmark} of {benchmark_pl.height} benchmark row(s); "
242 f"{joined_dates.height} row(s) remain. Pass a benchmark covering "
243 f"the same dates as the returns to avoid this.",
244 BenchmarkAlignmentWarning,
245 stacklevel=2,
246 )
247 returns_pl = returns_pl.join(joined_dates, on=date_col, how="inner")
248 benchmark_pl = benchmark_pl.join(joined_dates, on=date_col, how="inner")
249 return returns_pl, benchmark_pl
252def _prices_to_returns(frame: pl.DataFrame, date_col: str, frame_name: str) -> pl.DataFrame:
253 """Convert a price-level frame to a returns frame via percentage change.
255 The first row is dropped because no prior price is available to compute a
256 return for it.
258 Args:
259 frame: Price-level frame with a *date_col* column and asset columns.
260 date_col: Name of the date column (passed through unchanged).
261 frame_name: Descriptive name used in the error message when *date_col*
262 is missing (e.g. ``"prices"`` or ``"benchmark"``).
264 Returns:
265 pl.DataFrame: Returns frame with the same columns as *frame*, one row
266 shorter.
268 Raises:
269 MissingDateColumnError: If *date_col* is not a column of *frame*.
270 """
271 if date_col not in frame.columns:
272 raise MissingDateColumnError(frame_name, column=date_col, available=list(frame.columns))
273 asset_cols = _value_columns(frame, date_col)
274 return frame.with_columns([pl.col(c).pct_change().alias(c) for c in asset_cols]).slice(1)