Coverage for src/jquantstats/_data_reshape.py: 100%

1"""Reshaping operations for `Data`: resampling, copying, slicing, truncation. 

2 

3`_ReshapeMixin` collects the methods that return a *new* `Data` derived from an 

4existing one. They are factored out of ``data.py`` to keep that module focused 

5on construction and accessors; the mixin only reads the three dataclass fields 

6(``returns``, ``index``, ``benchmark``) and rebuilds via `_rebuild`, which 

7constructs through ``type(self)`` so `Data` never enters this module's import 

8graph — not even lazily. 

9""" 

10 

11from __future__ import annotations 

12 

13from datetime import date, datetime 

14from typing import TYPE_CHECKING, cast 

15 

16import polars as pl 

17 

18from ._truncate import resolve_bounds 

19 

20if TYPE_CHECKING: 

21 from collections.abc import Callable 

22 

23 from .data import Data 

24 

25 

26class _ReshapeMixin: 

27 """Mixin providing the `Data` operations that yield a new `Data`. 

28 

29 The concrete class (`Data`) supplies the ``returns``, ``index`` and 

30 ``benchmark`` dataclass fields; they are annotated here so the mixin's 

31 methods type-check without importing `Data` at module load (which would 

32 re-form an import cycle). No runtime attributes are created — the mixin 

33 carries empty slots. 

34 """ 

35 

36 __slots__ = () 

37 

38 # Provided by the concrete Data dataclass; declared for type-checkers only. 

39 returns: pl.DataFrame 

40 index: pl.DataFrame 

41 benchmark: pl.DataFrame | None 

42 

43 def _rebuild( 

44 self, 

45 *, 

46 returns: pl.DataFrame, 

47 index: pl.DataFrame, 

48 benchmark: pl.DataFrame | None = None, 

49 ) -> Data: 

50 """Build a fresh `Data` from the given frames. 

51 

52 Constructs via ``type(self)`` rather than importing `Data`. This mixin is 

53 only ever mixed into `Data`, so ``type(self)`` *is* the concrete class at 

54 runtime — which keeps `Data` out of this module's import graph entirely 

55 (a lazy import still puts it there) and rebuilds a subclass as its own 

56 type rather than downcasting it to `Data`. 

57 

58 Args: 

59 returns: Returns frame for the new object. 

60 index: Date/row index frame for the new object. 

61 benchmark: Optional benchmark frame for the new object. 

62 

63 Returns: 

64 Data: A new `Data` built from the supplied frames. 

65 """ 

66 factory = cast("Callable[..., Data]", type(self)) 

67 return factory(returns=returns, index=index, benchmark=benchmark) 

68 

69 def resample(self, every: str = "1mo") -> Data: 

70 """Resample returns and benchmark to a different frequency. 

71 

72 Args: 

73 every (str): Resampling frequency (e.g., ``'1mo'``, ``'1y'``). 

74 Defaults to ``'1mo'``. 

75 

76 Returns: 

77 Data: Resampled data at the requested frequency. 

78 

79 """ 

80 

81 def resample_frame(dframe: pl.DataFrame) -> pl.DataFrame: 

82 """Resample a single DataFrame to the target frequency using compound returns.""" 

83 dframe = self.index.hstack(dframe) # Add the date column for resampling 

84 

85 return dframe.group_by_dynamic( 

86 index_column=self.index.columns[0], every=every, period=every, closed="right", label="right" 

87 ).agg( 

88 [ 

89 ((pl.col(col) + 1.0).product() - 1.0).alias(col) 

90 for col in dframe.columns 

91 if col != self.index.columns[0] 

92 ] 

93 ) 

94 

95 resampled_returns = resample_frame(self.returns) 

96 resampled_benchmark = resample_frame(self.benchmark) if self.benchmark is not None else None 

97 resampled_index = resampled_returns.select(self.index.columns[0]) 

98 

99 return self._rebuild( 

100 returns=resampled_returns.drop(self.index.columns[0]), 

101 benchmark=resampled_benchmark.drop(self.index.columns[0]) if resampled_benchmark is not None else None, 

102 index=resampled_index, 

103 ) 

104 

105 def copy(self) -> Data: 

106 """Create a deep copy of the Data object. 

107 

108 Returns: 

109 Data: A new Data object with copies of the returns and benchmark. 

110 

111 """ 

112 benchmark = self.benchmark.clone() if self.benchmark is not None else None 

113 return self._rebuild(returns=self.returns.clone(), benchmark=benchmark, index=self.index.clone()) 

114 

115 def head(self, n: int = 5) -> Data: 

116 """Return the first n rows of the combined returns and benchmark data. 

117 

118 Args: 

119 n (int, optional): Number of rows to return. Defaults to 5. 

120 

121 Returns: 

122 Data: A new Data object containing the first n rows of the combined data. 

123 

124 """ 

125 benchmark_head = self.benchmark.head(n) if self.benchmark is not None else None 

126 return self._rebuild(returns=self.returns.head(n), benchmark=benchmark_head, index=self.index.head(n)) 

127 

128 def tail(self, n: int = 5) -> Data: 

129 """Return the last n rows of the combined returns and benchmark data. 

130 

131 Args: 

132 n (int, optional): Number of rows to return. Defaults to 5. 

133 

134 Returns: 

135 Data: A new Data object containing the last n rows of the combined data. 

136 

137 """ 

138 benchmark_tail = self.benchmark.tail(n) if self.benchmark is not None else None 

139 return self._rebuild(returns=self.returns.tail(n), benchmark=benchmark_tail, index=self.index.tail(n)) 

140 

141 def truncate( 

142 self, 

143 start: date | datetime | str | int | None = None, 

144 end: date | datetime | str | int | None = None, 

145 ) -> Data: 

146 """Return a new Data object truncated to the inclusive [start, end] range. 

147 

148 **The bound type picks the axis.** A ``date``, ``datetime`` or ISO-8601 

149 string is compared against the date column; an ``int`` is a 0-based row 

150 index and slices positionally. Row indices work on a temporal index too 

151 — ``truncate(start=10)`` drops the first ten rows of dated data — but the 

152 two kinds cannot be combined in one call. 

153 

154 An integer-indexed Data (no temporal index) accepts row indices only. 

155 

156 Args: 

157 start: Optional inclusive lower bound. A ``date``/``datetime``, an 

158 ISO-8601 string, or an ``int`` row index; ``int`` only when the 

159 index is not temporal. 

160 end: Optional inclusive upper bound. Same type rules as ``start``, 

161 and must address the same axis. 

162 

163 Returns: 

164 Data: A new Data object filtered to the specified range. 

165 

166 Raises: 

167 IntegerIndexBoundError: When the index is not temporal and a bound 

168 is not an ``int``. 

169 InvalidTruncateBoundError: When a bound is of an unsupported type, 

170 or is a string that is not ISO-8601. 

171 MixedTruncateBoundsError: When one bound is a row index and the 

172 other a date. 

173 """ 

174 date_column = self.index.columns[0] 

175 mode, lower, upper = resolve_bounds(start, end, temporal=self.index[date_column].dtype.is_temporal()) 

176 

177 if mode == "dates": 

178 new_index, new_returns, new_benchmark = self._truncate_temporal(date_column, lower, upper) 

179 else: 

180 # "none" resolves to a full-width slice, so it needs no separate branch. 

181 # The casts record what resolve_bounds guarantees for these modes but 

182 # cannot express in its return type. 

183 new_index, new_returns, new_benchmark = self._truncate_integer( 

184 cast("int | None", lower), cast("int | None", upper) 

185 ) 

186 

187 return self._rebuild(returns=new_returns, benchmark=new_benchmark, index=new_index) 

188 

189 def _truncate_temporal( 

190 self, 

191 date_column: str, 

192 start: date | datetime | str | int | None, 

193 end: date | datetime | str | int | None, 

194 ) -> tuple[pl.DataFrame, pl.DataFrame, pl.DataFrame | None]: 

195 """Truncate a temporal index by comparing the date column to [start, end].""" 

196 cond = pl.lit(True) 

197 if start is not None: 

198 cond = cond & (pl.col(date_column) >= pl.lit(start)) 

199 if end is not None: 

200 cond = cond & (pl.col(date_column) <= pl.lit(end)) 

201 mask = self.index.select(cond.alias("mask"))["mask"] 

202 new_benchmark = self.benchmark.filter(mask) if self.benchmark is not None else None 

203 return self.index.filter(mask), self.returns.filter(mask), new_benchmark 

204 

205 def _truncate_integer( 

206 self, 

207 start: int | None, 

208 end: int | None, 

209 ) -> tuple[pl.DataFrame, pl.DataFrame, pl.DataFrame | None]: 

210 """Truncate by 0-based row slicing. 

211 

212 Bounds arrive already validated by `resolve_bounds`, so this only 

213 substitutes the open-ended defaults. ``length`` is clamped at 0 so an 

214 inverted range yields an empty frame rather than a negative slice. 

215 """ 

216 row_start = start if start is not None else 0 

217 row_end = (end if end is not None else self.index.height - 1) + 1 

218 length = max(0, row_end - row_start) 

219 new_benchmark = self.benchmark.slice(row_start, length) if self.benchmark is not None else None 

220 return self.index.slice(row_start, length), self.returns.slice(row_start, length), new_benchmark