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 .exceptions import IntegerIndexBoundError 

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 When the index is temporal (Date/Datetime), truncation is performed by 

149 comparing the date column against ``start`` and ``end`` values. 

150 

151 When the index is integer-based, row slicing is used instead, and 

152 ``start`` and ``end`` must be non-negative integers. Passing 

153 non-integer bounds to an integer-indexed Data raises `TypeError`. 

154 

155 Args: 

156 start: Optional lower bound (inclusive). A date/datetime value 

157 when the index is temporal; a non-negative `int` row 

158 index when the data has no temporal index. 

159 end: Optional upper bound (inclusive). Same type rules as 

160 ``start``. 

161 

162 Returns: 

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

164 

165 Raises: 

166 TypeError: When the index is not temporal and a non-integer bound 

167 is supplied. 

168 

169 """ 

170 date_column = self.index.columns[0] 

171 

172 if self.index[date_column].dtype.is_temporal(): 

173 new_index, new_returns, new_benchmark = self._truncate_temporal(date_column, start, end) 

174 else: 

175 new_index, new_returns, new_benchmark = self._truncate_integer(start, end) 

176 

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

178 

179 def _truncate_temporal( 

180 self, 

181 date_column: str, 

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

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

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

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

186 cond = pl.lit(True) 

187 if start is not None: 

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

189 if end is not None: 

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

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

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

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

194 

195 @staticmethod 

196 def _resolve_row_bound(name: str, value: date | datetime | str | int | None, default: int) -> int: 

197 """Validate and resolve an integer truncation bound to a row index. 

198 

199 Args: 

200 name: The bound's name (``"start"`` or ``"end"``) for the message. 

201 value: The supplied bound; ``None`` and ``int`` are accepted. 

202 default: The row index to use when *value* is ``None``. 

203 

204 Returns: 

205 int: *value* when it is an ``int``, otherwise *default*. 

206 

207 Raises: 

208 IntegerIndexBoundError: If *value* is neither ``None`` nor ``int``. 

209 """ 

210 if value is not None and not isinstance(value, int): 

211 raise IntegerIndexBoundError(name, type(value).__name__) 

212 return value if value is not None else default 

213 

214 def _truncate_integer( 

215 self, 

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

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

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

219 """Truncate an integer index by row slicing; bounds must be integers.""" 

220 row_start = self._resolve_row_bound("start", start, 0) 

221 row_end = self._resolve_row_bound("end", end, self.index.height - 1) + 1 

222 length = max(0, row_end - row_start) 

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

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