feat: adds Web3SubscriptionsManager.get_subscription_data_nowait()

silverback/subscriptions.py

@@ -1,7 +1,7 @@
 import asyncio
 import json
 from enum import Enum
-from typing import AsyncGenerator
+from typing import AsyncGenerator, Union
 
 from ape.logging import logger
 from websockets import ConnectionClosedError
@@ -115,6 +115,9 @@ async def subscribe(self, type: SubscriptionType, **filter_params) -> str:
         return sub_id
 
     async def get_subscription_data(self, sub_id: str) -> AsyncGenerator[dict, None]:
+        """Iterate items from the subscription queue. If nothing is in the
+        queue, await.
+        """
         while True:
             if not (queue := self._subscriptions.get(sub_id)) or queue.empty():
                 async with self._ws_lock:
@@ -124,6 +127,39 @@ async def get_subscription_data(self, sub_id: str) -> AsyncGenerator[dict, None]
             else:
                 yield await queue.get()
 
+    async def get_subscription_data_nowait(self, sub_id: str) -> AsyncGenerator[dict, None]:
+        """Iterate items from the subscription queue. If nothing is in the
+        queue, return.
+        """
+        while True:
+            if not (queue := self._subscriptions.get(sub_id)) or queue.empty():
+                async with self._ws_lock:
+                    # Keep pulling until a message comes to process
+                    # NOTE: Python <3.10 does not support `anext` function
+                    await self.__anext__()
+            else:
+                try:
+                    yield queue.get_nowait()
+                except asyncio.QueueEmpty:
+                    pass
+
+    async def pop_subscription_data(self, sub_id: str) -> Union[dict, None]:
+        """Remove and return a single item from the subscription queue."""
+
+        async with self._ws_lock:
+            # NOTE: Python <3.10 does not support `anext` function
+            await self.__anext__()
+
+        queue = self._subscriptions.get(sub_id)
+
+        if queue:
+            try:
+                return queue.get_nowait()
+            except asyncio.QueueEmpty:
+                pass
+
+        return None
+
     async def unsubscribe(self, sub_id: str) -> bool:
         if sub_id not in self._subscriptions:
             raise ValueError(f"Unknown sub_id '{sub_id}'")