Coverage for src/jquantstats/_utils/_construction.py: 100%

1"""Frame-construction and interpolation helpers for building `Data` objects. 

2 

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""" 

9 

10from __future__ import annotations 

11 

12import warnings 

13from typing import Literal 

14 

15import narwhals as nw 

16import polars as pl 

17 

18from .._types import NativeFrame 

19from ..exceptions import ( 

20 BenchmarkAlignmentWarning, 

21 MissingDateColumnError, 

22 NullsInReturnsError, 

23) 

24 

25__all__ = [ 

26 "interpolate", 

27] 

28 

29#: The canonical name of the date column on every `Data` object, matching the 

30#: name the `Portfolio` internals already enforce. Inputs are renamed to it once 

31#: at construction so that ``data.index.columns[0]`` is the same string however 

32#: the object was built. 

33DATE_COLUMN = "date" 

34 

35 

36def _to_polars(df: NativeFrame) -> pl.DataFrame: 

37 """Convert any narwhals-compatible DataFrame to a polars DataFrame.""" 

38 if isinstance(df, pl.DataFrame): 

39 return df 

40 return nw.from_native(df, eager_only=True).to_polars() 

41 

42 

43def _value_columns(dframe: pl.DataFrame, date_col: str) -> list[str]: 

44 """Return every column of *dframe* except the date column.""" 

45 return [c for c in dframe.columns if c != date_col] 

46 

47 

48def _columns_with_nulls(dframe: pl.DataFrame, value_cols: list[str]) -> list[str]: 

49 """Return the subset of *value_cols* that contain at least one null.""" 

50 null_counts = dframe.select(value_cols).null_count().row(0) 

51 return [col for col, count in zip(value_cols, null_counts, strict=False) if count > 0] 

52 

53 

54def _apply_null_strategy( 

55 dframe: pl.DataFrame, 

56 date_col: str, 

57 frame_name: str, 

58 null_strategy: Literal["raise", "drop", "forward_fill"] | None, 

59) -> pl.DataFrame: 

60 """Check for nulls in *dframe* and apply *null_strategy*. 

61 

62 Args: 

63 dframe (pl.DataFrame): DataFrame to inspect. The date column is 

64 excluded from the null scan. 

65 date_col (str): Name of the column to treat as the date index 

66 (excluded from null check). 

67 frame_name (str): Descriptive name used in the error message 

68 (e.g. ``"returns"``). 

69 null_strategy ({"raise", "drop", "forward_fill"} | None): How to 

70 handle null values: 

71 

72 - ``None`` — leave nulls as-is (nulls will propagate through 

73 calculations). 

74 - ``"raise"`` — raise `NullsInReturnsError` if any null is found. 

75 - ``"drop"`` — drop every row that contains at least one null. 

76 - ``"forward_fill"`` — fill each null with the most recent 

77 non-null value in the same column. 

78 

79 Returns: 

80 pl.DataFrame: The original DataFrame (``None`` / ``"raise"``), a 

81 filtered DataFrame (``"drop"``), or a filled DataFrame 

82 (``"forward_fill"``). 

83 

84 Raises: 

85 NullsInReturnsError: When *null_strategy* is ``"raise"`` and nulls 

86 are present. 

87 

88 """ 

89 if null_strategy is None: 

90 return dframe 

91 

92 value_cols = _value_columns(dframe, date_col) 

93 cols_with_nulls = _columns_with_nulls(dframe, value_cols) 

94 

95 if not cols_with_nulls: 

96 return dframe 

97 

98 if null_strategy == "raise": 

99 raise NullsInReturnsError(frame_name, cols_with_nulls) 

100 if null_strategy == "drop": 

101 return dframe.drop_nulls(subset=value_cols) 

102 # forward_fill 

103 return dframe.with_columns(pl.col(value_cols).forward_fill()) 

104 

105 

106def interpolate(df: pl.DataFrame) -> pl.DataFrame: 

107 """Forward-fill numeric columns only between first and last non-null values. 

108 

109 For each numeric column, forward-fill is applied strictly within the span 

110 bounded by its first and last non-null samples. Values outside this span 

111 are left as-is (including leading/trailing nulls). Non-numeric columns are 

112 returned unchanged. 

113 

114 Args: 

115 df: Input frame possibly containing nulls. 

116 

117 Returns: 

118 pl.DataFrame: Frame where numeric columns have been interior-forward- 

119 filled; schema and dtypes of the original columns are preserved. 

120 

121 Examples: 

122 >>> import polars as pl 

123 >>> from jquantstats import interpolate 

124 >>> df = pl.DataFrame({"a": [None, 1.0, None, 3.0, None], "b": ["x", "y", "z", "w", "v"]}) 

125 >>> result = interpolate(df) 

126 

127 The interior null is filled; the leading and trailing ones are left alone: 

128 

129 >>> result["a"].to_list() 

130 [None, 1.0, 1.0, 3.0, None] 

131 

132 Non-numeric columns pass through untouched: 

133 

134 >>> result["b"].to_list() 

135 ['x', 'y', 'z', 'w', 'v'] 

136 

137 """ 

138 # Choose a temp column name guaranteed not to collide with any user column. 

139 tmp_col = "__row_idx__" 

140 while tmp_col in df.columns: 

141 tmp_col = f"_{tmp_col}_" 

142 

143 out = [] 

144 

145 for col in df.columns: 

146 s = df[col] 

147 if s.dtype.is_numeric(): 

148 non_null_mask = s.is_not_null() 

149 if non_null_mask.any(): 

150 _fwd = non_null_mask.arg_max() 

151 _rev = non_null_mask.reverse().arg_max() 

152 if _fwd is None or _rev is None: # pragma: no cover 

153 out.append(pl.col(col)) 

154 continue 

155 first_valid_idx = _fwd 

156 last_valid_idx = len(s) - 1 - _rev 

157 else: 

158 out.append(pl.col(col)) 

159 continue 

160 

161 mask = (pl.col(tmp_col) >= pl.lit(first_valid_idx)) & (pl.col(tmp_col) <= pl.lit(last_valid_idx)) 

162 filled_col = pl.when(mask).then(pl.col(col).fill_null(strategy="forward")).otherwise(pl.col(col)).alias(col) 

163 out.append(filled_col) 

164 else: 

165 out.append(pl.col(col)) 

166 

167 return df.with_columns(pl.int_range(0, df.height).alias(tmp_col)).select(out) 

168 

169 

170def _subtract_risk_free(dframe: pl.DataFrame, rf: float | pl.DataFrame, date_col: str) -> pl.DataFrame: 

171 """Subtract the risk-free rate from all numeric columns in the DataFrame. 

172 

173 Args: 

174 dframe (pl.DataFrame): DataFrame containing returns data with a date 

175 column and one or more numeric columns representing asset returns. 

176 rf (float | pl.DataFrame): Risk-free rate to subtract from returns. 

177 

178 - If float: A constant risk-free rate applied to all dates. 

179 - If pl.DataFrame: A DataFrame with a date column and a second 

180 column containing time-varying risk-free rates. 

181 

182 date_col (str): Name of the date column in both DataFrames for 

183 joining when rf is a DataFrame. 

184 

185 Returns: 

186 pl.DataFrame: DataFrame with the risk-free rate subtracted from all 

187 numeric columns, preserving the original column names. 

188 

189 """ 

190 if isinstance(rf, float): 

191 rf_dframe = dframe.select([pl.col(date_col), pl.lit(rf).alias("rf")]) 

192 else: 

193 if not isinstance(rf, pl.DataFrame): 

194 raise TypeError("rf must be a float or DataFrame") # noqa: TRY003 

195 if rf.columns[1] != "rf": 

196 warnings.warn( 

197 f"Risk-free rate column '{rf.columns[1]}' has been renamed to 'rf' for internal alignment.", 

198 stacklevel=3, 

199 ) 

200 rf_dframe = rf.rename({rf.columns[1]: "rf"}) if rf.columns[1] != "rf" else rf 

201 

202 dframe = dframe.join(rf_dframe, on=date_col, how="inner") 

203 return dframe.select( 

204 [pl.col(date_col)] 

205 + [(pl.col(col) - pl.col("rf")).alias(col) for col in dframe.columns if col not in {date_col, "rf"}] 

206 ) 

207 

208 

209def _canonicalise_date_column(frame: pl.DataFrame, frame_name: str, date_col: str | None) -> tuple[pl.DataFrame, str]: 

210 """Resolve *frame*'s date column and rename a temporal one to `DATE_COLUMN`. 

211 

212 When *date_col* is ``None`` the date column is auto-detected as the first 

213 temporal column of *frame*, matching what `Portfolio` does at construction. 

214 An explicit *date_col* is used verbatim and need not be temporal, so an 

215 integer or string index column can still be nominated by name — such a 

216 column keeps its own name, since calling an integer axis ``'date'`` would 

217 be a lie. 

218 

219 Args: 

220 frame: A returns, prices, benchmark or risk-free frame. 

221 frame_name: Descriptive name used in the error message (e.g. ``"returns"``). 

222 date_col: Name of the date column, or ``None`` to auto-detect it. 

223 

224 Returns: 

225 tuple[pl.DataFrame, str]: The frame and the name its date column now 

226 carries — ``'date'`` for a temporal axis, the nominated name otherwise. 

227 

228 Raises: 

229 MissingDateColumnError: If *date_col* is not a column of *frame*, or — 

230 when auto-detecting — if *frame* has no temporal column at all. 

231 """ 

232 if date_col is None: 

233 temporal = [name for name, dtype in frame.schema.items() if dtype.is_temporal()] 

234 if not temporal: 

235 raise MissingDateColumnError(frame_name, available=list(frame.columns)) 

236 source = temporal[0] 

237 elif date_col not in frame.columns: 

238 raise MissingDateColumnError(frame_name, column=date_col, available=list(frame.columns)) 

239 else: 

240 source = date_col 

241 

242 if source == DATE_COLUMN or not frame.schema[source].is_temporal(): 

243 return frame, source 

244 return frame.rename({source: DATE_COLUMN}), DATE_COLUMN 

245 

246 

247def _align_returns_benchmark( 

248 returns_pl: pl.DataFrame, benchmark_pl: pl.DataFrame, date_col: str 

249) -> tuple[pl.DataFrame, pl.DataFrame]: 

250 """Inner-join returns and benchmark on their common dates. 

251 

252 Args: 

253 returns_pl: Returns frame with a *date_col* column. 

254 benchmark_pl: Benchmark frame with a *date_col* column. 

255 date_col: The shared date column name. 

256 

257 Returns: 

258 The two frames filtered to their overlapping dates. 

259 

260 Raises: 

261 ValueError: If the frames share no dates. 

262 

263 Warns: 

264 BenchmarkAlignmentWarning: If aligning drops rows from either frame. 

265 """ 

266 joined_dates = returns_pl.join(benchmark_pl, on=date_col, how="inner").select(date_col) 

267 if joined_dates.is_empty(): 

268 raise ValueError("No overlapping dates between returns and benchmark.") # noqa: TRY003 

269 dropped_returns = returns_pl.height - joined_dates.height 

270 dropped_benchmark = benchmark_pl.height - joined_dates.height 

271 if dropped_returns > 0 or dropped_benchmark > 0: 

272 warnings.warn( 

273 f"Aligning returns and benchmark on common dates dropped " 

274 f"{dropped_returns} of {returns_pl.height} returns row(s) and " 

275 f"{dropped_benchmark} of {benchmark_pl.height} benchmark row(s); " 

276 f"{joined_dates.height} row(s) remain. Pass a benchmark covering " 

277 f"the same dates as the returns to avoid this.", 

278 BenchmarkAlignmentWarning, 

279 stacklevel=2, 

280 ) 

281 returns_pl = returns_pl.join(joined_dates, on=date_col, how="inner") 

282 benchmark_pl = benchmark_pl.join(joined_dates, on=date_col, how="inner") 

283 return returns_pl, benchmark_pl 

284 

285 

286def _prices_to_returns(frame: pl.DataFrame, date_col: str) -> pl.DataFrame: 

287 """Convert a price-level frame to a returns frame via percentage change. 

288 

289 The first row is dropped because no prior price is available to compute a 

290 return for it. 

291 

292 Args: 

293 frame: Price-level frame with a *date_col* column and asset columns. 

294 Callers canonicalise the date column first, so its presence is a 

295 precondition here rather than something re-checked. 

296 date_col: Name of the date column (passed through unchanged). 

297 

298 Returns: 

299 pl.DataFrame: Returns frame with the same columns as *frame*, one row 

300 shorter. 

301 """ 

302 asset_cols = _value_columns(frame, date_col) 

303 return frame.with_columns([pl.col(c).pct_change().alias(c) for c in asset_cols]).slice(1)