Skip to content

Commit

Permalink
feat: type annotate async function (and an internal function)
Browse files Browse the repository at this point in the history
This adds basic type annotation to async_stream_zip (and another internal
function). And there is a bit of rejigging to get types to check, but other
than adding return types the client-facing API should be the same.

This is inspired by #120

There will probably be more changes after this PR.
  • Loading branch information
michalc committed May 25, 2024
1 parent d9d79b4 commit 5abae39
Showing 1 changed file with 8 additions and 7 deletions.
15 changes: 8 additions & 7 deletions stream_zip.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
import asyncio
import secrets
import zlib
from typing import Any, Iterable, Tuple, Optional, Deque, Type, AsyncIterable
from typing import Any, Iterable, Tuple, Optional, Deque, Type, AsyncIterable, Callable

from Crypto.Cipher import AES
from Crypto.Hash import HMAC, SHA1
Expand Down Expand Up @@ -156,7 +156,7 @@ def _raise_if_beyond(offset: int, maximum: int, exception_class: Type[Exception]
if offset > maximum:
raise exception_class()

def _with_returned(gen):
def _with_returned(gen) -> Tuple[Callable[[], Optional[Any]], Iterable[bytes]]:
# We leverage the not-often used "return value" of generators. Here, we want to iterate
# over chunks (to encrypt them), but still return the same "return value". So we use a
# bit of a trick to extract the return value but still have access to the chunks as
Expand Down Expand Up @@ -719,7 +719,7 @@ def _no_compression_streamed_data(chunks, uncompressed_size, crc_32, maximum_siz

async def async_stream_zip(member_files, *args, **kwargs) -> AsyncIterable[bytes]:

async def to_async_iterable(sync_iterable):
async def to_async_iterable(sync_iterable) -> AsyncIterable[Any]:
# asyncio.to_thread is not available until Python 3.9, and StopIteration doesn't get
# propagated by run_in_executor, so we use a sentinel to detect the end of the iterable
done = object()
Expand All @@ -729,17 +729,18 @@ async def to_async_iterable(sync_iterable):
try:
import contextvars
except ImportError:
get_args = lambda: (next, it, done)
get_func_args: Callable[[], Tuple[Callable[..., Any], Tuple[Any, ...]]] = lambda: (next, (it, done))
else:
get_args = lambda: (contextvars.copy_context().run, next, it, done)
get_func_args = lambda: (contextvars.copy_context().run, (next, it, done))

while True:
value = await loop.run_in_executor(None, *get_args())
func, args = get_func_args()
value = await loop.run_in_executor(None, func, *args)
if value is done:
break
yield value

def to_sync_iterable(async_iterable):
def to_sync_iterable(async_iterable) -> Iterable[Any]:
# The built-in aiter and anext functions are not available until Python 3.10
async_it = async_iterable.__aiter__()
while True:
Expand Down

0 comments on commit 5abae39

Please sign in to comment.