BufferReader.range(join=True): trim and join chunks (#14)

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.
 

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.

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):

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 = []

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)