diff --git a/a_sync/primitives/queue.py b/a_sync/primitives/queue.py index bb8458cb..bfbcf3bc 100644 --- a/a_sync/primitives/queue.py +++ b/a_sync/primitives/queue.py @@ -1,6 +1,6 @@ import asyncio import sys -from typing import Generic, TypeVar +from typing import Generic, List, TypeVar, Union, overload T = TypeVar('T') @@ -11,11 +11,19 @@ class Queue(*bases): """The only difference between an a_sync.Queue and an asyncio.Queue is that `get_nowait` can retrn multiple responses.""" - def get_nowait(self, i: int = 1, can_return_less: bool = False) -> T: + @overload + def get_nowait(self, i = 1, can_return_less: bool = False) -> T: + ... + @overload + def get_nowait(self, i: int, can_return_less: bool = False) -> List[T]: + ... + def get_nowait(self, i: int = 1, can_return_less: bool = False) -> Union[T, List[T]]: """ Just like `asyncio.Queue.get_nowait`, but will return `i` items instead of 1. Set `can_return_less` to True if you want to receive up to `i` items. """ + if can_return_less and i == 1: + raise ValueError("you cant set i == 1 with can_return_less == True") values = [] if i == -1: while True: @@ -30,4 +38,4 @@ def get_nowait(self, i: int = 1, can_return_less: bool = False) -> T: if can_return_less: break raise - return values \ No newline at end of file + return values[0] if i == 1 else values \ No newline at end of file