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

1"""Shared bound resolution for `Data.truncate` and `Portfolio.truncate`. 

2 

3Both entry points accept the same bound types and must agree on what they mean, 

4so the classification lives here rather than being written twice. 

5 

6The rule is: **the bound type picks the axis, the index type says what is 

7legal.** An ``int`` bound is a 0-based row index, a ``date``/``datetime``/ISO-8601 

8string is a position on the temporal axis, and the two cannot be mixed. An 

9object with an integer index accepts row indices only. 

10 

11Routing on the bound rather than on the index dtype is what makes an ``int`` 

12bound meaningful on dated data: the previous behaviour compared the row number 

13against the date column, which Polars evaluated to an all-true mask, so 

14``truncate(start=10)`` silently returned every row. 

15""" 

16 

17from __future__ import annotations 

18 

19from datetime import date, datetime 

20from typing import Any, Literal 

21 

22from .exceptions import IntegerIndexBoundError, InvalidTruncateBoundError, MixedTruncateBoundsError 

23 

24#: What the resolved bounds address: nothing, row indices, or the temporal axis. 

25Mode = Literal["none", "rows", "dates"] 

26 

27 

28def _parse_iso(param: str, value: str) -> date | datetime: 

29 """Parse an ISO-8601 date or datetime string. 

30 

31 Args: 

32 param: Name of the parameter being parsed, for the error message. 

33 value: The candidate string. 

34 

35 Returns: 

36 A `datetime.date` for a plain date, a `datetime.datetime` when the 

37 string carries a time component. 

38 

39 Raises: 

40 InvalidTruncateBoundError: If *value* is not ISO-8601. 

41 """ 

42 try: 

43 return date.fromisoformat(value) 

44 except ValueError: 

45 pass 

46 try: 

47 return datetime.fromisoformat(value) 

48 except ValueError: 

49 raise InvalidTruncateBoundError(param, value) from None 

50 

51 

52def _classify(param: str, value: Any) -> tuple[Mode, Any]: 

53 """Classify one bound and coerce it to the form the comparison needs. 

54 

55 ``bool`` is rejected rather than accepted as an ``int``: ``truncate(start=True)`` 

56 is far more likely to be a mistake than a request for row 1. 

57 

58 Args: 

59 param: Name of the parameter, for the error message. 

60 value: The supplied bound. 

61 

62 Returns: 

63 A ``(mode, coerced)`` pair; mode is ``"none"`` when *value* is ``None``. 

64 

65 Raises: 

66 InvalidTruncateBoundError: If *value* is of an unsupported type, or is a 

67 string that is not ISO-8601. 

68 """ 

69 if value is None: 

70 return "none", None 

71 if isinstance(value, bool): 

72 raise InvalidTruncateBoundError(param, value) 

73 if isinstance(value, int): 

74 return "rows", value 

75 if isinstance(value, (date, datetime)): 

76 return "dates", value 

77 if isinstance(value, str): 

78 return "dates", _parse_iso(param, value) 

79 raise InvalidTruncateBoundError(param, value) 

80 

81 

82def _resolve_row_bounds(start: Any, end: Any) -> tuple[Mode, Any, Any]: 

83 """Resolve bounds against an integer index, where only row numbers are legal. 

84 

85 Split out of `resolve_bounds` so the non-temporal rule reads as one thing: 

86 a row number is the only meaning a bound can carry here, so anything else is 

87 reported against that expectation rather than being classified further. 

88 

89 ``bool`` is rejected for the reason `_classify` rejects it — ``start=True`` 

90 is far more likely to be a mistake than a request for row 1 — but the error 

91 differs: on an integer index the complaint is the index type, not the bound 

92 type, so `IntegerIndexBoundError` names what was expected. 

93 

94 Args: 

95 start: Inclusive lower bound, or ``None``. 

96 end: Inclusive upper bound, or ``None``. 

97 

98 Returns: 

99 ``(mode, start, end)`` with the bounds unchanged — a row index needs no 

100 coercion. ``mode`` is ``"none"`` when both bounds are ``None``, else 

101 ``"rows"``. 

102 

103 Raises: 

104 IntegerIndexBoundError: If either bound is not an ``int``. 

105 """ 

106 for param, value in (("start", start), ("end", end)): 

107 if value is not None and (isinstance(value, bool) or not isinstance(value, int)): 

108 raise IntegerIndexBoundError(param, type(value).__name__) 

109 return ("none" if start is None and end is None else "rows"), start, end 

110 

111 

112def _reject_mixed_modes(start_mode: Mode, end_mode: Mode) -> None: 

113 """Reject a row index paired with a date, in either order. 

114 

115 The two checks are mirror images, and the argument order to 

116 `MixedTruncateBoundsError` is not symmetric: the first name is the bound 

117 that disagrees with the axis already established by the other, so the 

118 message reads as a complaint about the newcomer rather than about the pair. 

119 

120 Args: 

121 start_mode: The mode `_classify` assigned to *start*. 

122 end_mode: The mode `_classify` assigned to *end*. 

123 

124 Raises: 

125 MixedTruncateBoundsError: If one bound is a row index and the other a date. 

126 """ 

127 if start_mode == "rows" and end_mode == "dates": 

128 raise MixedTruncateBoundsError("start", "end") 

129 if start_mode == "dates" and end_mode == "rows": 

130 raise MixedTruncateBoundsError("end", "start") 

131 

132 

133def resolve_bounds(start: Any, end: Any, *, temporal: bool) -> tuple[Mode, Any, Any]: 

134 """Decide which axis *start* and *end* address, and coerce them for it. 

135 

136 Args: 

137 start: Inclusive lower bound, or ``None``. 

138 end: Inclusive upper bound, or ``None``. 

139 temporal: Whether the object has a temporal index. When ``False`` the 

140 only legal bound is an ``int`` row index. 

141 

142 Returns: 

143 ``(mode, start, end)``. ``mode`` is ``"none"`` when both bounds are 

144 ``None`` (the caller should return the object unchanged), ``"rows"`` for 

145 positional slicing, or ``"dates"`` for a comparison against the date 

146 column. The returned bounds are coerced — ISO strings become 

147 `datetime.date` or `datetime.datetime` values. 

148 

149 Raises: 

150 IntegerIndexBoundError: If the index is not temporal and a bound is not 

151 an ``int``. 

152 InvalidTruncateBoundError: If a bound is of an unsupported type, or is a 

153 string that is not ISO-8601. 

154 MixedTruncateBoundsError: If one bound is a row index and the other a date. 

155 """ 

156 if not temporal: 

157 return _resolve_row_bounds(start, end) 

158 

159 start_mode, start_value = _classify("start", start) 

160 end_mode, end_value = _classify("end", end) 

161 _reject_mixed_modes(start_mode, end_mode) 

162 

163 # With the modes agreed, whichever bound is set names the axis; when only one 

164 # is given the other stays "none" and contributes nothing. 

165 mode: Mode = start_mode if start_mode != "none" else end_mode 

166 return mode, start_value, end_value