Fix closeable scoped/borrowed iter (#69)

* explicitly testing closing of scoped_iter and repeated slicing

* __aiter__ of a borrowed iterator no longer exposes original iterator

* fixed doc typos

* asyncstdlib related type annotations

* turned on mypy strict checking for asynctools and _core (see #65)

Co-authored-by: isra17 <[email protected]>

Author: maxfischer2781

asyncstdlib/asynctools.py

@@ -27,34 +27,27 @@ class _BorrowedAsyncIterator(AsyncGenerator[T, S]):
     Borrowed async iterator/generator, preventing to ``aclose`` the ``iterable``
     """
 
-    # adding special methods such as `__aiter__` as `__slots__` allows to set them
+    # adding special methods such as `__anext__` as `__slots__` allows to set them
     # on the instance: the interpreter expects *descriptors* not methods, and
     # `__slots__` are descriptors just like methods.
-    __slots__ = "__wrapped__", "__aiter__", "__anext__", "asend", "athrow"
+    __slots__ = "__wrapped__", "__anext__", "asend", "athrow", "_wrapper"
 
     # Type checker does not understand `__slot__` definitions
-    __aiter__: Callable[[Any], AsyncGenerator[T, S]]
     __anext__: Callable[[Any], Awaitable[T]]
     asend: Any
     athrow: Any
 
     def __init__(self, iterator: Union[AsyncIterator[T], AsyncGenerator[T, S]]):
         self.__wrapped__ = iterator
-        # iterator.__aiter__ is likely to return iterator (e.g. for async def: yield)
-        # We wrap it in a separate async iterator/generator to hide its __aiter__.
-        try:
-            wrapped_iterator: AsyncGenerator[T, S] = self._wrapped_iterator(iterator)
-            self.__anext__ = iterator.__anext__  # type: ignore
-            self.__aiter__ = wrapped_iterator.__aiter__  # type: ignore
-        except (AttributeError, TypeError):
-            raise TypeError(
-                "borrowing requires an async iterator "
-                + f"with __aiter__ and __anext__ method, got {type(iterator).__name__}"
-            ) from None
-        self.__anext__ = wrapped_iterator.__anext__  # type: ignore
-        # Our wrapper cannot pass on asend/athrow without getting much heavier.
-        # Since interleaving anext/asend/athrow is not allowed, and the wrapper holds
-        # no internal state other than the iterator, circumventing it should be fine.
+        # Create an actual async generator wrapper that we can close. Otherwise,
+        # if we pass on the original iterator methods we cannot disable them if
+        # anyone has a reference to them.
+        self._wrapper: AsyncGenerator[T, S] = self._wrapped_iterator(iterator)
+        # Forward all async iterator/generator methods but __aiter__ and aclose:
+        # An async *iterator* (e.g. `async def: yield`) must return
+        # itself from __aiter__. If we do not shadow this then
+        # running aiter(self).aclose closes the underlying iterator.
+        self.__anext__ = self._wrapper.__anext__  # type: ignore
         if hasattr(iterator, "asend"):
             self.asend = iterator.asend  # type: ignore
         if hasattr(iterator, "athrow"):
@@ -70,11 +63,14 @@ async def _wrapped_iterator(
         async for item in iterator:
             yield item
 
-    def __repr__(self):
+    def __aiter__(self) -> AsyncGenerator[T, S]:
+        return self
+
+    def __repr__(self) -> str:
         return f"<asyncstdlib.borrow of {self.__wrapped__!r} at 0x{(id(self)):x}>"
 
-    async def _aclose_wrapper(self):
-        wrapper_iterator = self.__aiter__()
+    async def _aclose_wrapper(self) -> None:
+        wrapper_iterator = self._wrapper
         # allow closing the intermediate wrapper
         # this prevents a resource warning if the wrapper is GC'd
         # the underlying iterator is NOT affected by this
@@ -85,17 +81,17 @@ async def _aclose_wrapper(self):
         if hasattr(self, "athrow"):
             self.athrow = wrapper_iterator.athrow
 
-    def aclose(self):
+    def aclose(self) -> Awaitable[None]:
         return self._aclose_wrapper()
 
 
 class _ScopedAsyncIterator(_BorrowedAsyncIterator[T, S]):
     __slots__ = ()
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"<asyncstdlib.scoped_iter of {self.__wrapped__!r} at 0x{(id(self)):x}>"
 
-    async def aclose(self):
+    async def aclose(self) -> None:
         pass
 
 
@@ -119,16 +115,16 @@ async def __aenter__(self) -> AsyncIterator[T]:
         borrowed_iter = self._borrowed_iter = _ScopedAsyncIterator(self._iterator)
         return borrowed_iter
 
-    async def __aexit__(self, exc_type, exc_val, exc_tb):
+    async def __aexit__(self, *args: Any) -> bool:
         await self._borrowed_iter._aclose_wrapper()  # type: ignore
         await self._iterator.aclose()  # type: ignore
         return False
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return f"<{self.__class__.__name__} of {self._iterator!r} at 0x{(id(self)):x}>"
 
 
-def borrow(iterator: AsyncIterator[T]) -> _BorrowedAsyncIterator[T, None]:
+def borrow(iterator: AsyncIterator[T]) -> AsyncIterator[T]:
     """
     Borrow an async iterator, preventing to ``aclose`` it
 
@@ -146,10 +142,15 @@ def borrow(iterator: AsyncIterator[T]) -> _BorrowedAsyncIterator[T, None]:
     .. seealso:: Use :py:func:`~.scoped_iter` to ensure an (async) iterable
                  is eventually closed and only :term:`borrowed <borrowing>` until then.
     """
+    if not hasattr(iterator, "__anext__") or not hasattr(iterator, "__aiter__"):
+        raise TypeError(
+            "borrowing requires an async iterator "
+            + f"with __aiter__ and __anext__ method, got {type(iterator).__name__}"
+        )
     return _BorrowedAsyncIterator(iterator)
 
 
-def scoped_iter(iterable: AnyIterable[T]):
+def scoped_iter(iterable: AnyIterable[T]) -> AsyncContextManager[AsyncIterator[T]]:
     """
     Context manager that provides an async iterator for an (async) ``iterable``
 
@@ -166,9 +167,9 @@ def scoped_iter(iterable: AnyIterable[T]):
         async def head_tail(iterable, leading=5, trailing=5):
             '''Provide the first ``leading`` and last ``trailing`` items'''
             # create async iterator valid for the entire block
-            async with scoped_iter(iterable) as async_iter:
+            async with a.scoped_iter(iterable) as async_iter:
                 # ... safely pass it on without it being closed ...
-                async for item in a.isclice(async_iter, leading):
+                async for item in a.islice(async_iter, leading):
                     yield item
                 tail = deque(maxlen=trailing)
                 # ... and use it again in the block
@@ -336,7 +337,7 @@ def sync(function: Callable[..., T]) -> Callable[..., Awaitable[T]]:
     ...
 
 
-def sync(function: Callable) -> Callable[..., Awaitable[T]]:
+def sync(function: Callable[..., T]) -> Callable[..., Any]:
     r"""
     Wraps a callable to ensure its result can be ``await``\ ed
 
@@ -372,10 +373,10 @@ async def main():
         return function
 
     @wraps(function)
-    async def async_wrapped(*args, **kwargs):
+    async def async_wrapped(*args: Any, **kwargs: Any) -> T:
         result = function(*args, **kwargs)
         if isinstance(result, Awaitable):
-            return await result
+            return await result  # type: ignore
         return result
 
     return async_wrapped

asyncstdlib/contextlib.py

@@ -179,11 +179,11 @@ def __init__(self, enter_result=None):
         self.enter_result = enter_result
 
     @overload
-    def __aenter__(self: "NullContext[None]") -> None:
+    async def __aenter__(self: "NullContext[None]") -> None:
         ...
 
     @overload
-    def __aenter__(self: "NullContext[T]") -> T:
+    async def __aenter__(self: "NullContext[T]") -> T:
         ...
 
     async def __aenter__(self):

pyproject.toml

@@ -47,3 +47,21 @@ no_implicit_optional = true
 warn_redundant_casts = true
 warn_unused_ignores = true
 warn_unreachable = true
+
+[[tool.mypy.overrides]]
+module = [
+    "asyncstdlib.asynctools",
+    "asyncstdlib._core",
+]
+disallow_any_generics = true
+disallow_subclassing_any = true
+disallow_untyped_calls = true
+disallow_untyped_defs = true
+disallow_incomplete_defs = true
+check_untyped_defs = true
+disallow_untyped_decorators = true
+no_implicit_optional = true
+warn_unused_ignores = true
+warn_return_any = true
+no_implicit_reexport = true
+strict_equality = true

unittests/test_asynctools.py

@@ -8,6 +8,23 @@
 CLOSED = "closed"
 
 
+@sync
+async def test_scoped_iter_iterclose():
+    """A `scoped_iter` cannot be closed via its public interface"""
+    async_iterable, iterable = asyncify(range(10)), iter(range(10))
+    async with a.scoped_iter(async_iterable) as a1:
+        assert await a.anext(a1) == next(iterable)
+        # closing a scoped iterator is a no-op
+        await a1.aclose()
+        assert await a.anext(a1) == next(iterable)
+        # explicitly test #68
+        await a.iter(a1).aclose()
+        assert await a.anext(a1) == next(iterable)
+        assert await a.list(async_iterable) == list(iterable)
+    assert await a.anext(a1, CLOSED) == CLOSED
+    assert await a.anext(async_iterable, CLOSED) == CLOSED
+
+
 @sync
 async def test_nested_lifetime():
     async_iterable, iterable = asyncify(range(10)), iter(range(10))

unittests/test_itertools.py

@@ -170,6 +170,16 @@ async def test_islice_exact(slicing):
     )
 
 
+@sync
+async def test_islice_scoped_iter():
+    """multiple `isclice` on borrowed iterator are consecutive"""
+    async_iterable, iterable = asyncify(range(10)), iter(range(10))
+    async with a.scoped_iter(async_iterable) as a1:
+        assert await a.list(a.islice(a1, 5)) == list(itertools.islice(iterable, 5))
+        assert await a.list(a.islice(a1, 5)) == list(itertools.islice(iterable, 5))
+        assert await a.list(a.islice(a1, 5)) == list(itertools.islice(iterable, 5))
+
+
 starmap_cases = [
     (lambda x, y: x + y, [(1, 2), (3, 4)]),
     (lambda *args: sum(args), [range(i) for i in range(1, 10)]),