Skip to content

Commit

Permalink
fix: Queue.get_nowait with i == 1 (#85)
Browse files Browse the repository at this point in the history
  • Loading branch information
BobTheBuidler authored Oct 11, 2023
1 parent 3c6d47c commit e887044
Showing 1 changed file with 11 additions and 3 deletions.
14 changes: 11 additions & 3 deletions a_sync/primitives/queue.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import asyncio
import sys
from typing import Generic, TypeVar
from typing import Generic, List, TypeVar, Union, overload

T = TypeVar('T')

Expand All @@ -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:
Expand All @@ -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
return values[0] if i == 1 else values

0 comments on commit e887044

Please sign in to comment.