Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# This file is part of Hypothesis, which may be found at
2# https://github.com/HypothesisWorks/hypothesis/
3#
4# Copyright the Hypothesis Authors.
5# Individual contributors are listed in AUTHORS.rst and the git log.
6#
7# This Source Code Form is subject to the terms of the Mozilla Public License,
8# v. 2.0. If a copy of the MPL was not distributed with this file, You can
9# obtain one at https://mozilla.org/MPL/2.0/.
11from collections.abc import Iterable, Iterator, Sequence
12from typing import Self, TypeAlias, cast, final
14IntervalsT: TypeAlias = tuple[tuple[int, int], ...]
17# @final makes mypy happy with the Self return annotations. We otherwise run
18# afoul of:
19# > You should not use Self as the return annotation if the method is not
20# > guaranteed to return an instance of a subclass when the class is subclassed
21# > https://docs.python.org/3/library/typing.html#typing.Self
24@final
25class IntervalSet:
26 """
27 A compact and efficient representation of a set of ``(a, b)`` intervals. Can
28 be treated like a set of integers, in that ``n in intervals`` will return
29 ``True`` if ``n`` is contained in any of the ``(a, b)`` intervals, and
30 ``False`` otherwise.
31 """
33 @classmethod
34 def from_string(cls, s: str) -> "Self":
35 """Return a tuple of intervals, covering the codepoints of characters in `s`.
37 >>> IntervalSet.from_string('abcdef0123456789')
38 ((48, 57), (97, 102))
39 """
40 x = cls([(ord(c), ord(c)) for c in sorted(s)])
41 return x.union(x)
43 def __init__(self, intervals: Iterable[Sequence[int]] = ()) -> None:
44 self.intervals: IntervalsT = cast(
45 IntervalsT, tuple(tuple(v) for v in intervals)
46 )
47 # cast above is validated by this length assertion. check here instead of
48 # before to not exhaust generators before we create intervals from it
49 assert all(len(v) == 2 for v in self.intervals)
51 self.offsets: list[int] = [0]
52 for u, v in self.intervals:
53 self.offsets.append(self.offsets[-1] + v - u + 1)
54 self.size = self.offsets.pop()
55 self._idx_of_zero = self.index_above(ord("0"))
56 self._idx_of_Z = min(self.index_above(ord("Z")), len(self) - 1)
58 def __len__(self) -> int:
59 return self.size
61 def __iter__(self) -> Iterator[int]:
62 for u, v in self.intervals:
63 yield from range(u, v + 1)
65 def __getitem__(self, i: int) -> int:
66 if i < 0:
67 i = self.size + i
68 if i < 0 or i >= self.size:
69 raise IndexError(f"Invalid index {i} for [0, {self.size})")
70 # Want j = maximal such that offsets[j] <= i
72 j = len(self.intervals) - 1
73 if self.offsets[j] > i:
74 hi = j
75 lo = 0
76 # Invariant: offsets[lo] <= i < offsets[hi]
77 while lo + 1 < hi:
78 mid = (lo + hi) // 2
79 if self.offsets[mid] <= i:
80 lo = mid
81 else:
82 hi = mid
83 j = lo
84 t = i - self.offsets[j]
85 u, v = self.intervals[j]
86 r = u + t
87 assert r <= v
88 return r
90 def __contains__(self, elem: str | int) -> bool:
91 if isinstance(elem, str):
92 elem = ord(elem)
93 assert 0 <= elem <= 0x10FFFF
94 return any(start <= elem <= end for start, end in self.intervals)
96 def __repr__(self) -> str:
97 return f"IntervalSet({self.intervals!r})"
99 def index(self, value: int) -> int:
100 for offset, (u, v) in zip(self.offsets, self.intervals, strict=True):
101 if u == value:
102 return offset
103 elif u > value:
104 raise ValueError(f"{value} is not in list")
105 if value <= v:
106 return offset + (value - u)
107 raise ValueError(f"{value} is not in list")
109 def index_above(self, value: int) -> int:
110 for offset, (u, v) in zip(self.offsets, self.intervals, strict=True):
111 if u >= value:
112 return offset
113 if value <= v:
114 return offset + (value - u)
115 return self.size
117 def __or__(self, other: "Self") -> "Self":
118 return self.union(other)
120 def __sub__(self, other: "Self") -> "Self":
121 return self.difference(other)
123 def __and__(self, other: "Self") -> "Self":
124 return self.intersection(other)
126 def __eq__(self, other: object) -> bool:
127 return isinstance(other, IntervalSet) and (other.intervals == self.intervals)
129 def __hash__(self) -> int:
130 return hash(self.intervals)
132 def union(self, other: "Self") -> "Self":
133 """Merge two sequences of intervals into a single tuple of intervals.
135 Any integer bounded by `x` or `y` is also bounded by the result.
137 >>> union([(3, 10)], [(1, 2), (5, 17)])
138 ((1, 17),)
139 """
140 assert isinstance(other, type(self))
141 x = self.intervals
142 y = other.intervals
143 if not x:
144 return IntervalSet(y)
145 if not y:
146 return IntervalSet(x)
147 intervals = sorted(x + y, reverse=True)
148 result = [intervals.pop()]
149 while intervals:
150 # 1. intervals is in descending order
151 # 2. pop() takes from the RHS.
152 # 3. (a, b) was popped 1st, then (u, v) was popped 2nd
153 # 4. Therefore: a <= u
154 # 5. We assume that u <= v and a <= b
155 # 6. So we need to handle 2 cases of overlap, and one disjoint case
156 # | u--v | u----v | u--v |
157 # | a----b | a--b | a--b |
158 u, v = intervals.pop()
159 a, b = result[-1]
160 if u <= b + 1:
161 # Overlap cases
162 result[-1] = (a, max(v, b))
163 else:
164 # Disjoint case
165 result.append((u, v))
166 return IntervalSet(result)
168 def difference(self, other: "Self") -> "Self":
169 """Set difference for lists of intervals. That is, returns a list of
170 intervals that bounds all values bounded by x that are not also bounded by
171 y. x and y are expected to be in sorted order.
173 For example difference([(1, 10)], [(2, 3), (9, 15)]) would
174 return [(1, 1), (4, 8)], removing the values 2, 3, 9 and 10 from the
175 interval.
176 """
177 assert isinstance(other, type(self))
178 x = self.intervals
179 y = other.intervals
180 if not y:
181 return IntervalSet(x)
182 x = list(map(list, x))
183 i = 0
184 j = 0
185 result: list[Iterable[int]] = []
186 while i < len(x) and j < len(y):
187 # Iterate in parallel over x and y. j stays pointing at the smallest
188 # interval in the left hand side that could still overlap with some
189 # element of x at index >= i.
190 # Similarly, i is not incremented until we know that it does not
191 # overlap with any element of y at index >= j.
193 xl, xr = x[i]
194 assert xl <= xr
195 yl, yr = y[j]
196 assert yl <= yr
198 if yr < xl:
199 # The interval at y[j] is strictly to the left of the interval at
200 # x[i], so will not overlap with it or any later interval of x.
201 j += 1
202 elif yl > xr:
203 # The interval at y[j] is strictly to the right of the interval at
204 # x[i], so all of x[i] goes into the result as no further intervals
205 # in y will intersect it.
206 result.append(x[i])
207 i += 1
208 elif yl <= xl:
209 if yr >= xr:
210 # x[i] is contained entirely in y[j], so we just skip over it
211 # without adding it to the result.
212 i += 1
213 else:
214 # The beginning of x[i] is contained in y[j], so we update the
215 # left endpoint of x[i] to remove this, and increment j as we
216 # now have moved past it. Note that this is not added to the
217 # result as is, as more intervals from y may intersect it so it
218 # may need updating further.
219 x[i][0] = yr + 1
220 j += 1
221 else:
222 # yl > xl, so the left hand part of x[i] is not contained in y[j],
223 # so there are some values we should add to the result.
224 result.append((xl, yl - 1))
226 if yr + 1 <= xr:
227 # If y[j] finishes before x[i] does, there may be some values
228 # in x[i] left that should go in the result (or they may be
229 # removed by a later interval in y), so we update x[i] to
230 # reflect that and increment j because it no longer overlaps
231 # with any remaining element of x.
232 x[i][0] = yr + 1
233 j += 1
234 else:
235 # Every element of x[i] other than the initial part we have
236 # already added is contained in y[j], so we move to the next
237 # interval.
238 i += 1
239 # Any remaining intervals in x do not overlap with any of y, as if they did
240 # we would not have incremented j to the end, so can be added to the result
241 # as they are.
242 result.extend(x[i:])
243 return IntervalSet(map(tuple, result))
245 def intersection(self, other: "Self") -> "Self":
246 """Set intersection for lists of intervals."""
247 assert isinstance(other, type(self)), other
248 intervals = []
249 i = j = 0
250 while i < len(self.intervals) and j < len(other.intervals):
251 u, v = self.intervals[i]
252 U, V = other.intervals[j]
253 if u > V:
254 j += 1
255 elif U > v:
256 i += 1
257 else:
258 intervals.append((max(u, U), min(v, V)))
259 if v < V:
260 i += 1
261 else:
262 j += 1
263 return IntervalSet(intervals)
265 def char_in_shrink_order(self, i: int) -> str:
266 # We would like it so that, where possible, shrinking replaces
267 # characters with simple ascii characters, so we rejig this
268 # bit so that the smallest values are 0, 1, 2, ..., Z.
269 #
270 # Imagine that numbers are laid out as abc0yyyZ...
271 # this rearranges them so that they are laid out as
272 # 0yyyZcba..., which gives a better shrinking order.
273 if i <= self._idx_of_Z:
274 # We want to rewrite the integers [0, n] inclusive
275 # to [zero_point, Z_point].
276 n = self._idx_of_Z - self._idx_of_zero
277 if i <= n:
278 i += self._idx_of_zero
279 else:
280 # We want to rewrite the integers [n + 1, Z_point] to
281 # [zero_point, 0] (reversing the order so that codepoints below
282 # zero_point shrink upwards).
283 i = self._idx_of_zero - (i - n)
284 assert i < self._idx_of_zero
285 assert 0 <= i <= self._idx_of_Z
287 return chr(self[i])
289 def index_from_char_in_shrink_order(self, c: str) -> int:
290 """
291 Inverse of char_in_shrink_order.
292 """
293 assert len(c) == 1
294 i = self.index(ord(c))
296 if i <= self._idx_of_Z:
297 n = self._idx_of_Z - self._idx_of_zero
298 # Rewrite [zero_point, Z_point] to [0, n].
299 if self._idx_of_zero <= i <= self._idx_of_Z:
300 i -= self._idx_of_zero
301 assert 0 <= i <= n
302 # Rewrite [zero_point, 0] to [n + 1, Z_point].
303 else:
304 i = self._idx_of_zero - i + n
305 assert n + 1 <= i <= self._idx_of_Z
306 assert 0 <= i <= self._idx_of_Z
308 return i