Merge branch 'fix_contextiterator' into 'main'

fix ContextIterator and release

See merge request !94

CHANGELOG.md

@@ -2,6 +2,12 @@
 
 ## (unreleased)
 
+## 0.5.1
+
+Bug fixes:
+
+- `ContextIterator` was not a proper generator API
+
 ## 0.5.0
 
 - `iter_bliss_data` and `iter_bliss_scan_data` are context iterators: context managers that yield an iterator

src/ewoksdata/__init__.py

-__version__ = "0.5.0"
+__version__ = "0.5.1"

src/ewoksdata/data/contextiterator.py

+import functools
+from typing import Iterator
 from collections.abc import Generator
 
 
 class ContextIterator(Generator):
-    def __init__(self, it):
+    """Can wrap an iterator or generator to ensure cleanup when used as a context manager."""
+
+    def __init__(self, it: Iterator):
         self._it = it
 
     def __iter__(self):
@@ -15,20 +19,24 @@ class ContextIterator(Generator):
         return self
 
     def __exit__(self, *_):
-        self._it.close()
+        if hasattr(self._it, "close"):
+            return self._it.close()
 
-    def send(self):
-        return next(self._it)
+    def send(self, value):
+        if hasattr(self._it, "send"):
+            return self._it.send(value)
 
-    def throw(self):
-        raise StopIteration
+    def throw(self, *args):
+        if hasattr(self._it, "throw"):
+            return self._it.throw(*args)
 
 
 def contextiterator(iterator):
-    """Decorator that allows to use an iterator as a context manager.
+    """Decorator that allows to use an iterator or generator as a context manager.
     This ensures that the iterator is closed when exiting the context manager.
     """
 
+    @functools.wraps(iterator)
     def wrapper(*args, **kw):
         return ContextIterator(iterator(*args, **kw))
 

src/ewoksdata/tests/test_contextiterator.py

+from ..data import contextiterator
+
+
+def test_context_iterator():
+
+    exited = False
+
+    @contextiterator.contextiterator
+    def iterator():
+        nonlocal exited
+        exited = False
+        try:
+            yield from [1, 2, 3]
+        finally:
+            exited = True
+
+    # Use like a normal iterator
+    it = iterator()
+    assert next(it) == 1
+    assert not exited
+
+    # Use like a context iterator
+    with iterator() as it:
+        assert next(it) == 1
+        assert not exited
+    assert exited
+
+
+def test_context_generator():
+
+    @contextiterator.contextiterator
+    def generator():
+        nonlocal exited
+        exited = False
+        try:
+            while True:
+                x = yield
+                yield x * 2
+        finally:
+            exited = True
+
+    # Use like a normal generator
+    exited = False
+    gen = generator()
+    assert next(gen) is None  # run up to first yield
+    assert not exited
+    assert gen.send(10) == 20
+    assert next(gen) is None
+    assert gen.send(15) == 30
+    assert not exited
+
+    # Use like a context iterator
+    exited = False
+    with generator() as gen:
+        assert next(gen) is None
+        assert not exited
+        assert gen.send(10) == 20
+        assert next(gen) is None
+        assert gen.send(15) == 30
+        assert not exited
+    assert exited