web3 7.0.0b2__py3-none-any.whl → 7.7.0__py3-none-any.whl

web3/providers/persistent/persistent_connection.py

@@ -2,9 +2,12 @@ from typing import (
     TYPE_CHECKING,
     Any,
     Dict,
+    Union,
+    cast,
 )
 
 from web3.types import (
+    FormattedEthSubscriptionResponse,
     RPCEndpoint,
     RPCResponse,
 )
@@ -16,6 +19,9 @@ if TYPE_CHECKING:
     from web3.manager import (  # noqa: F401
         _AsyncPersistentMessageStream,
     )
+    from web3.providers.persistent import (  # noqa: F401
+        PersistentConnectionProvider,
+    )
 
 
 class PersistentConnection:
@@ -26,17 +32,60 @@ class PersistentConnection:
 
     def __init__(self, w3: "AsyncWeb3"):
         self._manager = w3.manager
+        self.provider = cast("PersistentConnectionProvider", self._manager.provider)
 
-    # -- public methods -- #
     @property
     def subscriptions(self) -> Dict[str, Any]:
+        """
+        Return the active subscriptions on the persistent connection.
+
+        :return: The active subscriptions on the persistent connection.
+        :rtype: Dict[str, Any]
+        """
         return self._manager._request_processor.active_subscriptions
 
-    async def send(self, method: RPCEndpoint, params: Any) -> RPCResponse:
-        return await self._manager.send(method, params)
+    async def make_request(self, method: RPCEndpoint, params: Any) -> RPCResponse:
+        """
+        Make a request to the persistent connection and return the response. This method
+        does not process the response as it would when invoking a method via the
+        appropriate module on the `AsyncWeb3` instance,
+        e.g. `w3.eth.get_block("latest")`.
+
+        :param method: The RPC method, e.g. `eth_getBlockByNumber`.
+        :param params: The RPC method parameters, e.g. `["0x1337", False]`.
+
+        :return: The unprocessed response from the persistent connection.
+        :rtype: RPCResponse
+        """
+        return await self.provider.make_request(method, params)
+
+    async def send(self, method: RPCEndpoint, params: Any) -> None:
+        """
+        Send a raw, unprocessed message to the persistent connection.
 
-    async def recv(self) -> Any:
-        return await self._manager._get_next_message()
+        :param method: The RPC method, e.g. `eth_getBlockByNumber`.
+        :param params: The RPC method parameters, e.g. `["0x1337", False]`.
+
+        :return: None
+        """
+        await self._manager.send(method, params)
+
+    async def recv(self) -> Union[RPCResponse, FormattedEthSubscriptionResponse]:
+        """
+        Receive the next unprocessed response for a request from the persistent
+        connection.
+
+        :return: The next unprocessed response for a request from the persistent
+                 connection.
+        :rtype: Union[RPCResponse, FormattedEthSubscriptionResponse]
+        """
+        return await self._manager.recv()
 
     def process_subscriptions(self) -> "_AsyncPersistentMessageStream":
+        """
+        Asynchronous iterator that yields messages from the subscription message stream.
+
+        :return: The subscription message stream.
+        :rtype: _AsyncPersistentMessageStream
+        """
         return self._manager._persistent_message_stream()

web3/providers/persistent/request_processor.py

@@ -1,22 +1,39 @@
 import asyncio
-from copy import (
-    copy,
-)
+import sys
 from typing import (
     TYPE_CHECKING,
     Any,
     Callable,
     Dict,
+    Generic,
     Optional,
     Tuple,
+    TypeVar,
+    Union,
+)
+
+from eth_utils.toolz import (
+    compose,
 )
 
+from web3._utils.batching import (
+    BATCH_REQUEST_ID,
+)
 from web3._utils.caching import (
     RequestInformation,
     generate_cache_key,
 )
+from web3.exceptions import (
+    SubscriptionProcessingFinished,
+    TaskNotRunning,
+    Web3ValueError,
+)
+from web3.providers.persistent.subscription_manager import (
+    SubscriptionContainer,
+)
 from web3.types import (
     RPCEndpoint,
+    RPCId,
     RPCResponse,
 )
 from web3.utils import (
@@ -28,22 +45,55 @@ if TYPE_CHECKING:
         PersistentConnectionProvider,
     )
 
+T = TypeVar("T")
+
+# TODO: This is an ugly hack for python 3.8. Remove this after we drop support for it
+#  and use `asyncio.Queue[T]` type directly in the `TaskReliantQueue` class.
+if sys.version_info >= (3, 9):
+
+    class _TaskReliantQueue(asyncio.Queue[T], Generic[T]):
+        pass
+
+else:
+
+    class _TaskReliantQueue(asyncio.Queue, Generic[T]):  # type: ignore
+        pass
+
+
+class TaskReliantQueue(_TaskReliantQueue[T]):
+    """
+    A queue that relies on a task to be running to process items in the queue.
+    """
+
+    async def get(self) -> T:
+        item = await super().get()
+        if isinstance(item, Exception):
+            # if the item is an exception, raise it so the task can handle this case
+            # more gracefully
+            raise item
+        return item
+
 
 class RequestProcessor:
     _subscription_queue_synced_with_ws_stream: bool = False
 
+    # set by the subscription manager when it is initialized
+    _subscription_container: Optional[SubscriptionContainer] = None
+
     def __init__(
         self,
         provider: "PersistentConnectionProvider",
         subscription_response_queue_size: int = 500,
     ) -> None:
         self._provider = provider
-
         self._request_information_cache: SimpleCache = SimpleCache(500)
         self._request_response_cache: SimpleCache = SimpleCache(500)
-        self._subscription_response_queue: asyncio.Queue[RPCResponse] = asyncio.Queue(
-            maxsize=subscription_response_queue_size
-        )
+        self._subscription_response_queue: TaskReliantQueue[
+            Union[RPCResponse, TaskNotRunning]
+        ] = TaskReliantQueue(maxsize=subscription_response_queue_size)
+        self._handler_subscription_queue: TaskReliantQueue[
+            Union[RPCResponse, TaskNotRunning, SubscriptionProcessingFinished]
+        ] = TaskReliantQueue(maxsize=subscription_response_queue_size)
 
     @property
     def active_subscriptions(self) -> Dict[str, Any]:
@@ -57,9 +107,14 @@ class RequestProcessor:
 
     def cache_request_information(
         self,
+        request_id: Optional[RPCId],
         method: RPCEndpoint,
         params: Any,
-        response_formatters: Tuple[Callable[..., Any], ...],
+        response_formatters: Tuple[
+            Union[Dict[str, Callable[..., Any]], Callable[..., Any]],
+            Callable[..., Any],
+            Callable[..., Any],
+        ],
     ) -> Optional[str]:
         cached_requests_key = generate_cache_key((method, params))
         if cached_requests_key in self._provider._request_cache._data:
@@ -73,13 +128,16 @@ class RequestProcessor:
                 )
                 return None
 
-        # copy the request counter and find the next request id without incrementing
-        # since this is done when / if the request is successfully sent
-        request_id = next(copy(self._provider.request_counter))
-        cache_key = generate_cache_key(request_id)
-
-        self._bump_cache_if_key_present(cache_key, request_id)
+        if request_id is None:
+            if not self._provider._is_batching:
+                raise Web3ValueError(
+                    "Request id must be provided when not batching requests."
+                )
+            # the _batch_request_counter is set when entering the context manager
+            request_id = self._provider._batch_request_counter
+            self._provider._batch_request_counter += 1
 
+        cache_key = generate_cache_key(request_id)
         request_info = RequestInformation(
             method,
             params,
@@ -95,30 +153,6 @@ class RequestProcessor:
         )
         return cache_key
 
-    def _bump_cache_if_key_present(self, cache_key: str, request_id: int) -> None:
-        """
-        If the cache key is present in the cache, bump the cache key and request id
-        by one to make room for the new request. This behavior is necessary when a
-        request is made but inner requests, say to `eth_estimateGas` if the `gas` is
-        missing, are made before the original request is sent.
-        """
-        if cache_key in self._request_information_cache:
-            original_request_info = self._request_information_cache.get_cache_entry(
-                cache_key
-            )
-            bump = generate_cache_key(request_id + 1)
-
-            # recursively bump the cache if the new key is also present
-            self._bump_cache_if_key_present(bump, request_id + 1)
-
-            self._provider.logger.debug(
-                "Caching internal request. Bumping original request in cache:\n"
-                f"    request_id=[{request_id}] -> [{request_id + 1}],\n"
-                f"    cache_key=[{cache_key}] -> [{bump}],\n"
-                f"    request_info={original_request_info.__dict__}"
-            )
-            self._request_information_cache.cache(bump, original_request_info)
-
     def pop_cached_request_information(
         self, cache_key: str
     ) -> Optional[RequestInformation]:
@@ -136,9 +170,9 @@ class RequestProcessor:
     ) -> RequestInformation:
         if "method" in response and response["method"] == "eth_subscription":
             if "params" not in response:
-                raise ValueError("Subscription response must have params field")
+                raise Web3ValueError("Subscription response must have params field")
             if "subscription" not in response["params"]:
-                raise ValueError(
+                raise Web3ValueError(
                     "Subscription response params must have subscription field"
                 )
 
@@ -150,9 +184,8 @@ class RequestProcessor:
                 # i.e. subscription request information remains in the cache
                 self._request_information_cache.get_cache_entry(cache_key)
             )
-
         else:
-            # retrieve the request info from the cache using the request id
+            # retrieve the request info from the cache using the response id
             cache_key = generate_cache_key(response["id"])
             if response in self._provider._request_cache._data.values():
                 request_info = (
@@ -181,6 +214,33 @@ class RequestProcessor:
 
         return request_info
 
+    def append_result_formatter_for_request(
+        self, request_id: int, result_formatter: Callable[..., Any]
+    ) -> None:
+        cache_key = generate_cache_key(request_id)
+        cached_request_info_for_id: RequestInformation = (
+            self._request_information_cache.get_cache_entry(cache_key)
+        )
+        if cached_request_info_for_id is not None:
+            (
+                current_result_formatters,
+                error_formatters,
+                null_result_formatters,
+            ) = cached_request_info_for_id.response_formatters
+            cached_request_info_for_id.response_formatters = (
+                compose(
+                    result_formatter,
+                    current_result_formatters,
+                ),
+                error_formatters,
+                null_result_formatters,
+            )
+        else:
+            self._provider.logger.debug(
+                f"No cached request info for response id `{request_id}`. Cannot "
+                f"append response formatter for response."
+            )
+
     def append_middleware_response_processor(
         self,
         response: RPCResponse,
@@ -215,7 +275,7 @@ class RequestProcessor:
     ) -> None:
         if subscription:
             if self._subscription_response_queue.full():
-                self._provider.logger.info(
+                self._provider.logger.debug(
                     "Subscription queue is full. Waiting for provider to consume "
                     "messages before caching."
                 )
@@ -225,7 +285,26 @@ class RequestProcessor:
             self._provider.logger.debug(
                 f"Caching subscription response:\n    response={raw_response}"
             )
-            await self._subscription_response_queue.put(raw_response)
+            subscription_id = raw_response.get("params", {}).get("subscription")
+            sub_container = self._subscription_container
+            if sub_container and sub_container.get_handler_subscription_by_id(
+                subscription_id
+            ):
+                # if the subscription has a handler, put it in the handler queue
+                await self._handler_subscription_queue.put(raw_response)
+            else:
+                # otherwise, put it in the subscription response queue so a response
+                # can be yielded by the message stream
+                await self._subscription_response_queue.put(raw_response)
+        elif isinstance(raw_response, list):
+            # Since only one batch should be in the cache at all times, we use a
+            # constant cache key for the batch response.
+            cache_key = generate_cache_key(BATCH_REQUEST_ID)
+            self._provider.logger.debug(
+                f"Caching batch response:\n    cache_key={cache_key},\n"
+                f"    response={raw_response}"
+            )
+            self._request_response_cache.cache(cache_key, raw_response)
         else:
             response_id = raw_response.get("id")
             cache_key = generate_cache_key(response_id)
@@ -235,19 +314,17 @@ class RequestProcessor:
             )
             self._request_response_cache.cache(cache_key, raw_response)
 
-    def pop_raw_response(
+    async def pop_raw_response(
         self, cache_key: str = None, subscription: bool = False
     ) -> Any:
         if subscription:
             qsize = self._subscription_response_queue.qsize()
-            if qsize == 0:
-                return None
+            raw_response = await self._subscription_response_queue.get()
 
             if not self._provider._listen_event.is_set():
                 self._provider._listen_event.set()
 
-            raw_response = self._subscription_response_queue.get_nowait()
-            if qsize == 1:
+            if qsize == 0:
                 if not self._subscription_queue_synced_with_ws_stream:
                     self._subscription_queue_synced_with_ws_stream = True
                     self._provider.logger.info(
@@ -268,7 +345,7 @@ class RequestProcessor:
             )
         else:
             if not cache_key:
-                raise ValueError(
+                raise Web3ValueError(
                     "Must provide cache key when popping a non-subscription response."
                 )
 
@@ -282,15 +359,18 @@ class RequestProcessor:
 
         return raw_response
 
-    # request processor class methods
+    # cache methods
 
-    def clear_caches(self) -> None:
-        """
-        Clear the request processor caches.
-        """
+    def _reset_handler_subscription_queue(self) -> None:
+        self._handler_subscription_queue = TaskReliantQueue(
+            maxsize=self._handler_subscription_queue.maxsize
+        )
 
+    def clear_caches(self) -> None:
+        """Clear the request processor caches."""
         self._request_information_cache.clear()
         self._request_response_cache.clear()
-        self._subscription_response_queue = asyncio.Queue(
+        self._subscription_response_queue = TaskReliantQueue(
             maxsize=self._subscription_response_queue.maxsize
         )
+        self._reset_handler_subscription_queue()

web3/providers/persistent/subscription_container.py

@@ -0,0 +1,56 @@
+from typing import (
+    Any,
+    Dict,
+    Iterator,
+    List,
+    Optional,
+)
+
+from eth_typing import (
+    HexStr,
+)
+
+from web3.utils import (
+    EthSubscription,
+)
+
+
+class SubscriptionContainer:
+    def __init__(self) -> None:
+        self.subscriptions: List[EthSubscription[Any]] = []
+        self.subscriptions_by_id: Dict[HexStr, EthSubscription[Any]] = {}
+        self.subscriptions_by_label: Dict[str, EthSubscription[Any]] = {}
+
+    def __len__(self) -> int:
+        return len(self.subscriptions)
+
+    def __iter__(self) -> Iterator[EthSubscription[Any]]:
+        return iter(self.subscriptions)
+
+    def add_subscription(self, subscription: EthSubscription[Any]) -> None:
+        self.subscriptions.append(subscription)
+        self.subscriptions_by_id[subscription.id] = subscription
+        self.subscriptions_by_label[subscription.label] = subscription
+
+    def remove_subscription(self, subscription: EthSubscription[Any]) -> None:
+        self.subscriptions.remove(subscription)
+        self.subscriptions_by_id.pop(subscription.id)
+        self.subscriptions_by_label.pop(subscription.label)
+
+    def get_by_id(self, sub_id: HexStr) -> EthSubscription[Any]:
+        return self.subscriptions_by_id.get(sub_id)
+
+    def get_by_label(self, label: str) -> EthSubscription[Any]:
+        return self.subscriptions_by_label.get(label)
+
+    @property
+    def handler_subscriptions(self) -> List[EthSubscription[Any]]:
+        return [sub for sub in self.subscriptions if sub._handler is not None]
+
+    def get_handler_subscription_by_id(
+        self, sub_id: HexStr
+    ) -> Optional[EthSubscription[Any]]:
+        sub = self.get_by_id(sub_id)
+        if sub and sub._handler:
+            return sub
+        return None

web3/providers/persistent/subscription_manager.py

@@ -0,0 +1,233 @@
+import asyncio
+import logging
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    List,
+    Sequence,
+    Union,
+    cast,
+    overload,
+)
+
+from eth_typing import (
+    HexStr,
+)
+
+from web3.exceptions import (
+    SubscriptionProcessingFinished,
+    TaskNotRunning,
+    Web3TypeError,
+    Web3ValueError,
+)
+from web3.providers.persistent.subscription_container import (
+    SubscriptionContainer,
+)
+from web3.types import (
+    FormattedEthSubscriptionResponse,
+    RPCResponse,
+)
+from web3.utils.subscriptions import (
+    EthSubscription,
+    EthSubscriptionContext,
+)
+
+if TYPE_CHECKING:
+    from web3 import AsyncWeb3  # noqa: F401
+    from web3.providers.persistent import (  # noqa: F401
+        PersistentConnectionProvider,
+        RequestProcessor,
+    )
+
+
+class SubscriptionManager:
+    """
+    The ``SubscriptionManager`` is responsible for subscribing, unsubscribing, and
+    managing all active subscriptions for an ``AsyncWeb3`` instance. It is also
+    used for processing all subscriptions that have handler functions.
+    """
+
+    logger: logging.Logger = logging.getLogger(
+        "web3.providers.persistent.subscription_manager"
+    )
+    total_handler_calls: int = 0
+
+    def __init__(self, w3: "AsyncWeb3") -> None:
+        self._w3 = w3
+        self._provider = cast("PersistentConnectionProvider", w3.provider)
+        self._subscription_container = SubscriptionContainer()
+
+        # share the subscription container with the request processor so it can separate
+        # subscriptions into different queues based on ``sub._handler`` presence
+        self._provider._request_processor._subscription_container = (
+            self._subscription_container
+        )
+
+    def _add_subscription(self, subscription: EthSubscription[Any]) -> None:
+        self._subscription_container.add_subscription(subscription)
+
+    def _remove_subscription(self, subscription: EthSubscription[Any]) -> None:
+        self._subscription_container.remove_subscription(subscription)
+
+    @property
+    def subscriptions(self) -> List[EthSubscription[Any]]:
+        return self._subscription_container.subscriptions
+
+    def get_by_id(self, sub_id: HexStr) -> EthSubscription[Any]:
+        return self._subscription_container.get_by_id(sub_id)
+
+    def get_by_label(self, label: str) -> EthSubscription[Any]:
+        return self._subscription_container.get_by_label(label)
+
+    @overload
+    async def subscribe(self, subscriptions: EthSubscription[Any]) -> HexStr:
+        ...
+
+    @overload
+    async def subscribe(
+        self, subscriptions: Sequence[EthSubscription[Any]]
+    ) -> List[HexStr]:
+        ...
+
+    async def subscribe(
+        self, subscriptions: Union[EthSubscription[Any], Sequence[EthSubscription[Any]]]
+    ) -> Union[HexStr, List[HexStr]]:
+        """
+        Used to subscribe to a single or multiple subscriptions.
+
+        :param subscriptions: A single subscription or a sequence of subscriptions.
+        :type subscriptions: Union[EthSubscription, Sequence[EthSubscription]]
+        :return:
+        """
+        if isinstance(subscriptions, EthSubscription):
+            if (
+                subscriptions.label
+                in self._subscription_container.subscriptions_by_label
+            ):
+                raise Web3ValueError(
+                    "Subscription label already exists. Subscriptions must have "
+                    f"unique labels.\n    label: {subscriptions.label}"
+                )
+
+            subscriptions.manager = self
+            sub_id = await self._w3.eth._subscribe(*subscriptions.subscription_params)
+            subscriptions._id = sub_id
+            self._add_subscription(subscriptions)
+            self.logger.info(
+                "Successfully subscribed to subscription:\n    "
+                f"label: {subscriptions.label}\n    id: {sub_id}"
+            )
+            return sub_id
+        elif isinstance(subscriptions, Sequence):
+            if len(subscriptions) == 0:
+                raise Web3ValueError("No subscriptions provided.")
+
+            sub_ids: List[HexStr] = []
+            for sub in subscriptions:
+                await self.subscribe(sub)
+            return sub_ids
+        raise Web3TypeError("Expected a Subscription or a sequence of Subscriptions.")
+
+    async def unsubscribe(self, subscription: EthSubscription[Any]) -> bool:
+        """
+        Used to unsubscribe from a subscription.
+
+        :param subscription: The subscription to unsubscribe from.
+        :type subscription: EthSubscription
+        :return: ``True`` if unsubscribing was successful, ``False`` otherwise.
+        :rtype: bool
+        """
+        if subscription not in self.subscriptions:
+            raise Web3ValueError(
+                "Subscription not found or is not being managed by the subscription "
+                f"manager.\n    label: {subscription.label}\n    id: {subscription._id}"
+            )
+        if await self._w3.eth._unsubscribe(subscription.id):
+            self._remove_subscription(subscription)
+            self.logger.info(
+                "Successfully unsubscribed from subscription:\n    "
+                f"label: {subscription.label}\n    id: {subscription.id}"
+            )
+            if len(self._subscription_container.handler_subscriptions) == 0:
+                queue = self._provider._request_processor._handler_subscription_queue
+                await queue.put(SubscriptionProcessingFinished())
+            return True
+        return False
+
+    async def unsubscribe_all(self) -> bool:
+        """
+        Used to unsubscribe from all subscriptions that are being managed by the
+        subscription manager.
+
+        :return: ``True`` if unsubscribing was successful, ``False`` otherwise.
+        :rtype: bool
+        """
+        unsubscribed = [
+            await self.unsubscribe(sub) for sub in self.subscriptions.copy()
+        ]
+        if all(unsubscribed):
+            self.logger.info("Successfully unsubscribed from all subscriptions.")
+            return True
+        else:
+            if len(self.subscriptions) > 0:
+                self.logger.warning(
+                    "Failed to unsubscribe from all subscriptions. Some subscriptions "
+                    f"are still active.\n    subscriptions={self.subscriptions}"
+                )
+            return False
+
+    async def handle_subscriptions(self, run_forever: bool = False) -> None:
+        """
+        Used to handle all subscriptions that have handlers. The method will run until
+        all subscriptions that have handler functions are unsubscribed from or, if
+        ``run_forever`` is set to ``True``, it will run indefinitely.
+
+        :param run_forever: If ``True``, the method will run indefinitely.
+        :type run_forever: bool
+        :return: None
+        """
+        if not self._subscription_container.handler_subscriptions:
+            self.logger.warning(
+                "No handler subscriptions found. Subscription handler did not run."
+            )
+            return
+
+        queue = self._provider._request_processor._handler_subscription_queue
+        try:
+            while run_forever or self._subscription_container.handler_subscriptions:
+                response = cast(RPCResponse, await queue.get())
+                formatted_sub_response = cast(
+                    FormattedEthSubscriptionResponse,
+                    await self._w3.manager._process_response(response),
+                )
+
+                # if the subscription was unsubscribed from, the response won't be
+                # formatted because we lost the request information
+                sub_id = formatted_sub_response.get("subscription")
+                sub = self._subscription_container.get_handler_subscription_by_id(
+                    sub_id
+                )
+                if sub:
+                    await sub._handler(
+                        EthSubscriptionContext(
+                            self._w3,
+                            sub,
+                            formatted_sub_response["result"],
+                            **sub._handler_context,
+                        )
+                    )
+        except SubscriptionProcessingFinished:
+            self.logger.info(
+                "All handler subscriptions have been unsubscribed from. "
+                "Stopping subscription handling."
+            )
+        except TaskNotRunning:
+            await asyncio.sleep(0)
+            self._provider._handle_listener_task_exceptions()
+            self.logger.error(
+                "Message listener background task for the provider has stopped "
+                "unexpectedly. Stopping subscription handling."
+            )
+
+        # no active handler subscriptions, clear the handler subscription queue
+        self._provider._request_processor._reset_handler_subscription_queue()