diff --git a/stream2py/buffer_reader.py b/stream2py/buffer_reader.py index 4d07d6f..eb52d43 100644 --- a/stream2py/buffer_reader.py +++ b/stream2py/buffer_reader.py @@ -122,6 +122,8 @@ def __init__( self._buffer = buffer self._last_item = None self._last_key = None + self._join_items = None # set by the buffer from the source reader, if offered + self._slice_item = None self._stop_event = stop_event self._sleep_time_on_iter_none_s = 0.1 self.read_size = read_size # read_size used by __next__ @@ -237,6 +239,7 @@ def range( only_new_items=False, start_le=False, stop_ge=False, + join=False, ): """Enables: 1. Get last n minutes @@ -255,8 +258,26 @@ def range( i.e. keys=[0, 10, 20], start=9 will include key=0 :param stop_ge: if True, raise ValueError when there is no key greater than or equal to stop in buffer, if ignore_no_item_found is also True, return None instead of ValueError - :return: list of items in range + :param join: if True, return ONE item instead of a list: the items in range + trimmed to exactly ``[start, stop]`` and joined, using the source reader's + ``slice_item`` and ``join_items`` (both must be implemented, otherwise + NotImplementedError is raised before anything is read). Implies + ``start_le=True`` so the chunk containing ``start`` is included, and is + incompatible with ``step``. The ``last_item`` cursor still moves to the + last *buffer* item, not the trimmed one. With ``only_new_items`` the + joined item starts after the cursor. When nothing is found the result is + the same as without ``join`` (``[]``, ``None`` or an exception). + :return: list of items in range (a single joined item if ``join``) """ + if join: + if step is not None: + raise ValueError("range(join=True) doesn't support step") + if self._join_items is None or self._slice_item is None: + raise NotImplementedError( + "range(join=True) needs the source reader to implement both " + "join_items(items) and slice_item(item, start, stop)" + ) + start_le = True if ignore_no_item_found is None: ignore_no_item_found = self._read_kwargs.get("ignore_no_item_found", False) with self._buffer.reader_lock() as reader: @@ -275,6 +296,7 @@ def range( _start = start if start > _next_key else _next_key else: _start = start + trim_start = _start # where the joined item must begin, before start_le widens if start_le is True: with suppress( ValueError @@ -299,8 +321,20 @@ def range( if ignore_no_item_found: return None raise e + if join and items: + return self._trim_and_join(items, trim_start, stop) return items + def _trim_and_join(self, items, start, stop): + """Trim the first and last of ``items`` to ``[start, stop]`` and join them. + + Only the ends can overhang the query (the keys are ordered), and a single + item is both ends at once.""" + items = list(items) + items[0] = self._slice_item(items[0], start, None) + items[-1] = self._slice_item(items[-1], None, stop) + return self._join_items(items) + def tail(self, *, peek=False, ignore_no_item_found=False, only_new_items=False): """Finds the last item in buffer. Raise ValueError if no item found. diff --git a/stream2py/source_reader.py b/stream2py/source_reader.py index 9aca226..f88833f 100644 --- a/stream2py/source_reader.py +++ b/stream2py/source_reader.py @@ -129,6 +129,32 @@ def key(self, data: Any) -> ComparableType: "Implement the 'key' method to convert data into a comparable value to sort by" ) + def join_items(self, items: list) -> Any: + """(Optional) Join consecutive data items into one data item. + + Implement together with ``slice_item`` to let ``BufferReader.range(..., + join=True)`` return the exact span asked for, as a single item (e.g. several + waveform chunks concatenated). Not abstract: sources whose items can't be + joined just don't override it. + + :param items: non-empty list of data items, in key order + :return: one data item + """ + raise NotImplementedError("Implement 'join_items' to support range(join=True)") + + def slice_item(self, item: Any, start=None, stop=None) -> Any: + """(Optional) The part of ``item`` whose keys fall within ``[start, stop]``. + + ``None`` for a bound means unbounded on that side; an item that needs no + trimming can be returned as is. See ``join_items``. Must be pure (items in the + buffer are shared, never mutate them) and keyed on absolute keys: it is applied + once per bound, so a lone item is sliced twice (``start`` then ``stop``). + + :param item: a data item as returned by ``read`` + :return: a data item of the same kind, possibly shorter + """ + raise NotImplementedError("Implement 'slice_item' to support range(join=True)") + @property def closed(self) -> bool: """True if the source reader is closed, False if it's open. diff --git a/stream2py/stream_buffer.py b/stream2py/stream_buffer.py index 902b4a0..d67ccba 100644 --- a/stream2py/stream_buffer.py +++ b/stream2py/stream_buffer.py @@ -24,6 +24,21 @@ DFLT_MAX_LEN = 10000 +def _overridden(source_reader, name): + """The bound method ``name`` if the source reader really implements it, else None. + + ``SourceReader`` defines ``join_items``/``slice_item`` as raising stubs, so + presence on the object doesn't mean support.""" + method = getattr(source_reader, name, None) + if method is None: + return None + from stream2py.source_reader import SourceReader + + if getattr(type(source_reader), name, None) is getattr(SourceReader, name, None): + return None + return method + + class _SourceBuffer: def __init__( self, @@ -33,6 +48,8 @@ def __init__( key=None, maxlen: int = 10000, buffer_reader_class: Type[BufferReader] = BufferReader, + join_items=None, + slice_item=None, ): """StreamBuffer helper class Thread safe buffer for reading and writing data items @@ -46,6 +63,8 @@ def __init__( self._buffer = RWLockSortedDeque([], key=key, maxlen=maxlen) self._source_reader_info = source_reader_info self.buffer_reader_class = buffer_reader_class + self._join_items = join_items + self._slice_item = slice_item def __len__(self): return len(self._buffer) @@ -67,16 +86,22 @@ def drop(self, n=1): writer.drop(n) def mk_reader(self, **read_kwargs): - return self.buffer_reader_class( + reader = self.buffer_reader_class( buffer=self._buffer, source_reader_info=self._source_reader_info, stop_event=self._stop_event, **read_kwargs, ) + # set after construction so custom buffer_reader_classes need not know about them + reader._join_items = self._join_items + reader._slice_item = self._slice_item + return reader def attach_reader(self, reader): reader._buffer = self._buffer reader._stop_event = self._stop_event + reader._join_items = self._join_items + reader._slice_item = self._slice_item class StreamBuffer: @@ -293,6 +318,8 @@ def _open(self): buffer_reader_class=getattr( self.source_reader, "buffer_reader_class", BufferReader ), + join_items=_overridden(self.source_reader, "join_items"), + slice_item=_overridden(self.source_reader, "slice_item"), ) # def _mk_contextualized_iterator(self): diff --git a/stream2py/tests/test_buffer_reader_stopiteration.py b/stream2py/tests/test_buffer_reader_stopiteration.py index 54c990e..163ff3a 100644 --- a/stream2py/tests/test_buffer_reader_stopiteration.py +++ b/stream2py/tests/test_buffer_reader_stopiteration.py @@ -39,7 +39,9 @@ def test_next_returns_none_when_no_data_but_running(): # Read all available data quickly import time - time.sleep(0.1) # Let some data accumulate + deadline = time.time() + 5 # wait for data rather than hope 0.1s is enough on a loaded runner + while not len(buffer.source_buffer) and time.time() < deadline: + time.sleep(0.01) # Read until we get None (no more data available yet) items = [] diff --git a/stream2py/tests/test_zz_range_join.py b/stream2py/tests/test_zz_range_join.py new file mode 100644 index 0000000..23a8216 --- /dev/null +++ b/stream2py/tests/test_zz_range_join.py @@ -0,0 +1,78 @@ +"""range(join=True) (i2mint/stream2py#14). + +Deliberately thread-free: the buffer and reader are driven directly, so these tests +can't compete for CPU with the timing-sensitive threaded tests around them.""" +import threading + +import pytest + +from stream2py import SourceReader +from stream2py.stream_buffer import _SourceBuffer, _overridden + + +class Chunks(SourceReader): + """Items are (first_key, [values...]); value i sits at key first_key + i.""" + + def open(self): + pass + + def read(self): + return None + + def close(self): + pass + + def key(self, data): + return data[0] + + def slice_item(self, item, start=None, stop=None): + first, values = item + lo = 0 if start is None else max(start - first, 0) + hi = len(values) if stop is None else stop - first + 1 + return (first + lo, values[lo:hi]) + + def join_items(self, items): + return (items[0][0], [v for _, vs in items for v in vs]) + + +class NoJoin(Chunks): + slice_item = SourceReader.slice_item + join_items = SourceReader.join_items + + +def _reader(source, n=5, size=4): + sb = _SourceBuffer( + {}, + threading.Event(), + key=source.key, + join_items=_overridden(source, "join_items"), + slice_item=_overridden(source, "slice_item"), + ) + for i in range(n): + sb.append((i * size, list(range(i * size, (i + 1) * size)))) + return sb.mk_reader() + + +def test_range_join_trims_and_joins_chunks(): + reader = _reader(Chunks()) + assert reader.range(6, 13, join=True) == (6, list(range(6, 14))) + assert reader.last_key == 12 # cursor stays on a real buffer item + assert reader.range(5, 5, join=True, peek=True) == (5, [5]) + assert reader.range(0, 19, join=True) == (0, list(range(20))) + # default behaviour is unchanged: whole chunks, from the first key >= start + assert [k for k, _ in reader.range(6, 13)] == [8, 12] + with pytest.raises(ValueError): + reader.range(0, 5, step=2, join=True) + + +def test_range_join_only_new_items_starts_after_cursor(): + reader = _reader(Chunks()) + reader.range(0, 5) # cursor now on the chunk with key 4 + assert reader.last_key == 4 + assert reader.range(0, 13, join=True, only_new_items=True)[0] == 8 + + +def test_range_join_needs_the_source_to_implement_it(): + assert _overridden(NoJoin(), "join_items") is None + with pytest.raises(NotImplementedError): + _reader(NoJoin()).range(0, 5, join=True)