Full strict typing (#73)

* strict typing for all sub-packages

* testing strict typing for all sub-packages

Author: maxfischer2781

asyncstdlib/_typing.py

@@ -53,6 +53,7 @@
 T4 = TypeVar("T4")
 T5 = TypeVar("T5")
 R = TypeVar("R", covariant=True)
+C = TypeVar("C", bound=Callable[..., Any])
 AC = TypeVar("AC", bound=Callable[..., Awaitable[Any]])
 
 #: Hashable Key
@@ -72,7 +73,7 @@ def __lt__(self: LT, other: LT) -> bool:
 
 
 class SupportsAdd(Protocol):
-    def __add__(self: ADD, other: ADD) -> bool:
+    def __add__(self: ADD, other: ADD) -> ADD:
         raise NotImplementedError
 
 

asyncstdlib/builtins.py

@@ -27,7 +27,19 @@
 __ANEXT_DEFAULT = Sentinel("<no default>")
 
 
-async def anext(iterator: AsyncIterator[T], default=__ANEXT_DEFAULT) -> T:
+@overload
+async def anext(iterator: AsyncIterator[T]) -> T:
+    ...
+
+
+@overload
+async def anext(iterator: AsyncIterator[T], default: T) -> T:
+    ...
+
+
+async def anext(
+    iterator: AsyncIterator[T], default: Union[Sentinel, T] = __ANEXT_DEFAULT
+) -> T:
     """
     Retrieve the next item from the async iterator
 
@@ -47,7 +59,7 @@ async def anext(iterator: AsyncIterator[T], default=__ANEXT_DEFAULT) -> T:
     except StopAsyncIteration:
         if default is __ANEXT_DEFAULT:
             raise
-        return default
+        return default  # type: ignore
 
 
 __ITER_DEFAULT = Sentinel("<no default>")
@@ -132,7 +144,7 @@ async def any(iterable: AnyIterable[T]) -> bool:
 def zip(
     __it1: AnyIterable[T1],
     *,
-    strict=False,
+    strict: bool = ...,
 ) -> AsyncIterator[Tuple[T1]]:
     ...
 
@@ -142,7 +154,7 @@ def zip(
     __it1: AnyIterable[T1],
     __it2: AnyIterable[T2],
     *,
-    strict=False,
+    strict: bool = ...,
 ) -> AsyncIterator[Tuple[T1, T2]]:
     ...
 
@@ -153,7 +165,7 @@ def zip(
     __it2: AnyIterable[T2],
     __it3: AnyIterable[T3],
     *,
-    strict=False,
+    strict: bool = ...,
 ) -> AsyncIterator[Tuple[T1, T2, T3]]:
     ...
 
@@ -165,7 +177,7 @@ def zip(
     __it3: AnyIterable[T3],
     __it4: AnyIterable[T4],
     *,
-    strict=False,
+    strict: bool = ...,
 ) -> AsyncIterator[Tuple[T1, T2, T3, T4]]:
     ...
 
@@ -178,7 +190,7 @@ def zip(
     __it4: AnyIterable[T4],
     __it5: AnyIterable[T5],
     *,
-    strict=False,
+    strict: bool = ...,
 ) -> AsyncIterator[Tuple[T1, T2, T3, T4, T5]]:
     ...
 
@@ -191,13 +203,13 @@ def zip(
     __it4: AnyIterable[Any],
     __it5: AnyIterable[Any],
     *iterables: AnyIterable[Any],
-    strict=False,
+    strict: bool = ...,
 ) -> AsyncIterator[Tuple[Any, ...]]:
     ...
 
 
 async def zip(
-    *iterables: AnyIterable[Any], strict=False
+    *iterables: AnyIterable[Any], strict: bool = False
 ) -> AsyncIterator[Tuple[Any, ...]]:
     """
     Create an async iterator that aggregates elements from each of the (async) iterables
@@ -240,15 +252,22 @@ async def zip(
                 await aclose()
 
 
-async def _zip_inner(aiters):
+async def _zip_inner(
+    aiters: Tuple[AsyncIterator[T], ...]
+) -> AsyncIterator[Tuple[T, ...]]:
+    """Direct zip transposing tuple-of-iterators to iterator-of-tuples"""
     try:
         while True:
             yield (*[await anext(it) for it in aiters],)
     except StopAsyncIteration:
         return
 
 
-async def _zip_inner_strict(aiters):
+async def _zip_inner_strict(
+    aiters: Tuple[AsyncIterator[T], ...]
+) -> AsyncIterator[Tuple[T, ...]]:
+    """Length aware zip checking that all iterators are equal length"""
+    # track index of the last iterator we tried to anext
     tried = 0
     try:
         while True:
@@ -541,13 +560,12 @@ async def _min_max(
     :param invert: compute ``max`` if ``True`` and ``min`` otherwise
     """
     async with ScopedIter(iterable) as item_iter:
-        best = await anext(item_iter, default=__MIN_MAX_DEFAULT)
-        if best is __MIN_MAX_DEFAULT:
-            if default is __MIN_MAX_DEFAULT:
-                name = "max" if invert else "min"
-                raise ValueError(f"{name}() arg is an empty sequence")
-            return default
-        if key is None:
+        best = await anext(item_iter, default=default)
+        # this implies that item_iter is empty and default is __MIN_MAX_DEFAULT
+        if best is __MIN_MAX_DEFAULT:  # type: ignore
+            name = "max" if invert else "min"
+            raise ValueError(f"{name}() arg is an empty sequence")
+        elif key is None:
             async for item in item_iter:
                 if invert ^ (item < best):
                     best = item
@@ -587,7 +605,9 @@ async def filter(
                     yield item
 
 
-async def enumerate(iterable: AnyIterable[T], start=0) -> AsyncIterator[Tuple[int, T]]:
+async def enumerate(
+    iterable: AnyIterable[T], start: int = 0
+) -> AsyncIterator[Tuple[int, T]]:
     """
     An async iterator of running count and element in an (async) iterable
 
@@ -685,11 +705,6 @@ async def set(iterable: Union[Iterable[T], AsyncIterable[T]] = ()) -> Set[T]:
     return {element async for element in aiter(iterable)}
 
 
-async def _identity(x: T) -> T:
-    """Asynchronous identity function, returns its argument unchanged"""
-    return x
-
-
 @overload
 async def sorted(
     iterable: AnyIterable[LT], *, key: None = ..., reverse: bool = ...
@@ -707,7 +722,7 @@ async def sorted(
 async def sorted(
     iterable: AnyIterable[T],
     *,
-    key: Optional[Callable[[T], Any]] = None,
+    key: Optional[Callable[[T], LT]] = None,
     reverse: bool = False,
 ) -> List[T]:
     """
@@ -730,11 +745,15 @@ async def sorted(
         It is guaranteed to be worst-case O(n log n) runtime.
     """
     if key is None:
+        # TODO: is this a worthwhile optimisation?
         try:
             return _sync_builtins.sorted(iterable, reverse=reverse)  # type: ignore
         except TypeError:
-            pass
-    key = _awaitify(key) if key is not None else _identity
-    keyed_items = [(await key(item), item) async for item in aiter(iterable)]
-    keyed_items.sort(key=lambda ki: ki[0], reverse=reverse)
-    return [item for key, item in keyed_items]
+            items = [item async for item in aiter(iterable)]
+            items.sort(reverse=reverse)
+            return items
+    else:
+        async_key = _awaitify(key)
+        keyed_items = [(await async_key(item), item) async for item in aiter(iterable)]
+        keyed_items.sort(key=lambda ki: ki[0], reverse=reverse)
+        return [item for key, item in keyed_items]

asyncstdlib/contextlib.py

@@ -15,18 +15,21 @@
 from functools import partial
 import sys
 
-from ._typing import Protocol, AsyncContextManager, ContextManager, T
+from ._typing import Protocol, AsyncContextManager, ContextManager, T, C
 from ._core import awaitify
 from ._utility import public_module, slot_get as _slot_get
 
 
+AnyContextManager = Union[AsyncContextManager[T], ContextManager[T]]
+
+
 # typing.AsyncContextManager uses contextlib.AbstractAsyncContextManager if available,
 # and a custom implementation otherwise. No need to replicate it.
 AbstractContextManager = AsyncContextManager
 
 
 class ACloseable(Protocol):
-    async def aclose(self):
+    async def aclose(self) -> None:
         """Asynchronously close this object"""
 
 
@@ -58,29 +61,31 @@ async def Context(*args, **kwargs):
     """
 
     @wraps(func)
-    def helper(*args, **kwds):
+    def helper(*args: Any, **kwds: Any) -> AsyncContextManager[T]:
         return _AsyncGeneratorContextManager(func, args, kwds)
 
     return helper
 
 
-class _AsyncGeneratorContextManager:
-    def __init__(self, func, args, kwds):
+class _AsyncGeneratorContextManager(Generic[T]):
+    def __init__(
+        self, func: Callable[..., AsyncGenerator[T, None]], args: Any, kwds: Any
+    ):
         self.gen = func(*args, **kwds)
         self.__doc__ = getattr(func, "__doc__", type(self).__doc__)
 
-    async def __aenter__(self):
+    async def __aenter__(self) -> T:
         try:
             return await self.gen.__anext__()
         except StopAsyncIteration:
             raise RuntimeError("generator did not yield to __aenter__") from None
 
-    async def __aexit__(self, exc_type, exc_val, exc_tb):
+    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> bool:
         if exc_type is None:
             try:
                 await self.gen.__anext__()
             except StopAsyncIteration:
-                return
+                return False
             else:
                 raise RuntimeError("generator did not stop after __aexit__")
         else:
@@ -99,6 +104,7 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
             except exc_type as exc:
                 if exc is not exc_val:
                     raise
+                return False
             else:
                 raise RuntimeError("generator did not stop after throw() in __aexit__")
 
@@ -134,8 +140,9 @@ def __init__(self, thing: AC):
     async def __aenter__(self) -> AC:
         return self.thing
 
-    async def __aexit__(self, exc_type, exc_val, exc_tb):
+    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> bool:
         await self.thing.aclose()
+        return False
 
 
 closing = Closing
@@ -175,7 +182,7 @@ def __init__(self: "NullContext[None]", enter_result: None = ...) -> None:
     def __init__(self: "NullContext[T]", enter_result: T) -> None:
         ...
 
-    def __init__(self, enter_result=None):
+    def __init__(self, enter_result: Optional[T] = None):
         self.enter_result = enter_result
 
     @overload
@@ -186,11 +193,11 @@ async def __aenter__(self: "NullContext[None]") -> None:
     async def __aenter__(self: "NullContext[T]") -> T:
         ...
 
-    async def __aenter__(self):
+    async def __aenter__(self) -> Optional[T]:
         return self.enter_result
 
-    async def __aexit__(self, exc_type, exc_val, exc_tb):
-        pass
+    async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> bool:
+        return False
 
 
 nullcontext = NullContext
@@ -199,8 +206,8 @@ async def __aexit__(self, exc_type, exc_val, exc_tb):
 SE = TypeVar(
     "SE",
     bound=Union[
-        AsyncContextManager,
-        ContextManager,
+        AsyncContextManager[Any],
+        ContextManager[Any],
         Callable[[Any, BaseException, Any], Optional[bool]],
         Callable[[Any, BaseException, Any], Awaitable[Optional[bool]]],
     ],
@@ -228,11 +235,13 @@ class ExitStack:
         There are no separate methods to distinguish async and regular arguments.
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         self._exit_callbacks: Deque[Callable[..., Awaitable[Optional[bool]]]] = deque()
 
     @staticmethod
-    async def _aexit_callback(callback, exc_type, exc_val, tb):
+    async def _aexit_callback(
+        callback: Callable[[], Awaitable[Any]], exc_type: Any, exc_val: Any, tb: Any
+    ) -> bool:
         """Invoke a callback as if it were an ``__aexit__`` method"""
         await callback()
         return False  # callbacks never suppress exceptions
@@ -298,7 +307,7 @@ def push(self, exit: SE) -> SE:
         self._exit_callbacks.append(aexit)
         return exit
 
-    def callback(self, callback: Callable, *args, **kwargs):
+    def callback(self, callback: C, *args: Any, **kwargs: Any) -> C:
         """
         Registers an arbitrary callback to be called with arguments on unwinding
 
@@ -312,11 +321,11 @@ def callback(self, callback: Callable, *args, **kwargs):
         This method does not change its argument, and can be used as a context manager.
         """
         self._exit_callbacks.append(
-            partial(self._aexit_callback, awaitify(partial(callback, *args, **kwargs)))
+            partial(self._aexit_callback, partial(awaitify(callback), *args, **kwargs))
         )
         return callback
 
-    async def enter_context(self, cm: AsyncContextManager):
+    async def enter_context(self, cm: AnyContextManager[T]) -> T:
         """
         Enter the supplied context manager, and register it for exit if successful
 
@@ -353,9 +362,9 @@ async def enter_context(self, cm: AsyncContextManager):
         else:
             context_value = await _slot_get(cm, "__aenter__")()
         self._exit_callbacks.append(aexit)
-        return context_value
+        return context_value  # type: ignore
 
-    async def aclose(self):
+    async def aclose(self) -> None:
         """
         Immediately unwind the context stack
 
@@ -371,7 +380,7 @@ def _stitch_context(
         exception: BaseException,
         context: BaseException,
         base_context: Optional[BaseException],
-    ):
+    ) -> None:
         """
         Emulate that `exception` was caused by an unhandled `context`
 
@@ -392,10 +401,10 @@ def _stitch_context(
         # we expect it to reference
         exception.__context__ = context
 
-    async def __aenter__(self):
+    async def __aenter__(self) -> "ExitStack":
         return self
 
-    async def __aexit__(self, exc_type, exc_val, tb):
+    async def __aexit__(self, exc_type: Any, exc_val: Any, tb: Any) -> bool:
         received_exc = exc_type is not None
         # Even if we don't handle an exception *right now*, we may be part
         # of an exception handler unwinding gracefully. This is our __context__.