web3 7.6.1__py3-none-any.whl → 7.8.0__py3-none-any.whl

web3/providers/persistent/request_processor.py

@@ -1,7 +1,4 @@
 import asyncio
-from copy import (
-    copy,
-)
 import sys
 from typing import (
     TYPE_CHECKING,
@@ -9,6 +6,7 @@ from typing import (
     Callable,
     Dict,
     Generic,
+    List,
     Optional,
     Tuple,
     TypeVar,
@@ -27,11 +25,16 @@ from web3._utils.caching import (
     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 (
@@ -75,18 +78,23 @@ class TaskReliantQueue(_TaskReliantQueue[T]):
 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: 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]:
@@ -100,6 +108,7 @@ class RequestProcessor:
 
     def cache_request_information(
         self,
+        request_id: Optional[RPCId],
         method: RPCEndpoint,
         params: Any,
         response_formatters: Tuple[
@@ -120,25 +129,23 @@ class RequestProcessor:
                 )
                 return None
 
-        if self._provider._is_batching:
+        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
-            current_request_id = self._provider._batch_request_counter
+            request_id = self._provider._batch_request_counter
             self._provider._batch_request_counter += 1
-        else:
-            # copy the request counter and find the next request id without incrementing
-            # since this is done when / if the request is successfully sent
-            current_request_id = next(copy(self._provider.request_counter))
-        cache_key = generate_cache_key(current_request_id)
-
-        self._bump_cache_if_key_present(cache_key, current_request_id)
 
+        cache_key = generate_cache_key(request_id)
         request_info = RequestInformation(
             method,
             params,
             response_formatters,
         )
         self._provider.logger.debug(
-            f"Caching request info:\n    request_id={current_request_id},\n"
+            f"Caching request info:\n    request_id={request_id},\n"
             f"    cache_key={cache_key},\n    request_info={request_info.__dict__}"
         )
         self._request_information_cache.cache(
@@ -147,30 +154,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]:
@@ -288,6 +271,15 @@ class RequestProcessor:
 
     # raw response cache
 
+    def _is_batch_response(
+        self, raw_response: Union[List[RPCResponse], RPCResponse]
+    ) -> bool:
+        return isinstance(raw_response, list) or (
+            isinstance(raw_response, dict)
+            and raw_response.get("id") is None
+            and self._provider._is_batching
+        )
+
     async def cache_raw_response(
         self, raw_response: Any, subscription: bool = False
     ) -> None:
@@ -303,8 +295,18 @@ class RequestProcessor:
             self._provider.logger.debug(
                 f"Caching subscription response:\n    response={raw_response}"
             )
-            await self._subscription_response_queue.put(raw_response)
-        elif isinstance(raw_response, list):
+            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 self._is_batch_response(raw_response):
             # 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)
@@ -367,7 +369,12 @@ class RequestProcessor:
 
         return raw_response
 
-    # request processor class methods
+    # cache methods
+
+    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."""
@@ -376,3 +383,4 @@ class RequestProcessor:
         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,298 @@
+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)
+
+    def _validate_and_normalize_label(self, subscription: EthSubscription[Any]) -> None:
+        if subscription.label == subscription._default_label:
+            # if no custom label was provided, generate a unique label
+            i = 2
+            while self.get_by_label(subscription._label) is not None:
+                subscription._label = f"{subscription._default_label}#{i}"
+                i += 1
+        else:
+            if (
+                subscription._label
+                in self._subscription_container.subscriptions_by_label
+            ):
+                raise Web3ValueError(
+                    "Subscription label already exists. Subscriptions must have unique "
+                    f"labels.\n    label: {subscription._label}"
+                )
+
+    @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):
+            subscriptions.manager = self
+            self._validate_and_normalize_label(subscriptions)
+            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:
+                sub_ids.append(await self.subscribe(sub))
+            return sub_ids
+        raise Web3TypeError("Expected a Subscription or a sequence of Subscriptions.")
+
+    @overload
+    async def unsubscribe(self, subscriptions: EthSubscription[Any]) -> bool:
+        ...
+
+    @overload
+    async def unsubscribe(self, subscriptions: HexStr) -> bool:
+        ...
+
+    @overload
+    async def unsubscribe(
+        self,
+        subscriptions: Sequence[Union[EthSubscription[Any], HexStr]],
+    ) -> bool:
+        ...
+
+    async def unsubscribe(
+        self,
+        subscriptions: Union[
+            EthSubscription[Any],
+            HexStr,
+            Sequence[Union[EthSubscription[Any], HexStr]],
+        ],
+    ) -> bool:
+        """
+        Used to unsubscribe from one or multiple subscriptions.
+
+        :param subscriptions: The subscription(s) to unsubscribe from.
+        :type subscriptions: Union[EthSubscription, Sequence[EthSubscription], HexStr,
+            Sequence[HexStr]]
+        :return: ``True`` if unsubscribing to all was successful, ``False`` otherwise
+            with a warning.
+        :rtype: bool
+        """
+        if isinstance(subscriptions, EthSubscription) or isinstance(subscriptions, str):
+            if isinstance(subscriptions, str):
+                subscription_id = subscriptions
+                subscriptions = self.get_by_id(subscription_id)
+                if subscriptions is None:
+                    raise Web3ValueError(
+                        "Subscription not found or is not being managed by the "
+                        f"subscription manager.\n    id: {subscription_id}"
+                    )
+
+            if subscriptions not in self.subscriptions:
+                raise Web3ValueError(
+                    "Subscription not found or is not being managed by the "
+                    "subscription manager.\n    "
+                    f"label: {subscriptions.label}\n    id: {subscriptions._id}"
+                )
+
+            if await self._w3.eth._unsubscribe(subscriptions.id):
+                self._remove_subscription(subscriptions)
+                self.logger.info(
+                    "Successfully unsubscribed from subscription:\n    "
+                    f"label: {subscriptions.label}\n    id: {subscriptions.id}"
+                )
+
+                if len(self._subscription_container.handler_subscriptions) == 0:
+                    queue = (
+                        self._provider._request_processor._handler_subscription_queue
+                    )
+                    await queue.put(SubscriptionProcessingFinished())
+                return True
+
+        elif isinstance(subscriptions, Sequence):
+            if len(subscriptions) == 0:
+                raise Web3ValueError("No subscriptions provided.")
+
+            unsubscribed: List[bool] = []
+            for sub in subscriptions:
+                if isinstance(sub, str):
+                    sub = HexStr(sub)
+                unsubscribed.append(await self.unsubscribe(sub))
+            return all(unsubscribed)
+
+        self.logger.warning(
+            f"Failed to unsubscribe from subscription\n    subscription={subscriptions}"
+        )
+        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 and not run_forever:
+            self.logger.warning(
+                "No handler subscriptions found. Subscription handler did not run."
+            )
+            return
+
+        queue = self._provider._request_processor._handler_subscription_queue
+        while run_forever or self._subscription_container.handler_subscriptions:
+            try:
+                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:
+                if not run_forever:
+                    self.logger.info(
+                        "All handler subscriptions have been unsubscribed from. "
+                        "Stopping subscription handling."
+                    )
+                    break
+            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()

web3/providers/persistent/websocket.py

@@ -59,8 +59,6 @@ class WebSocketProvider(PersistentConnectionProvider):
     logger = logging.getLogger("web3.providers.WebSocketProvider")
     is_async: bool = True
 
-    _ws: Optional[WebSocketClientProtocol] = None
-
     def __init__(
         self,
         endpoint_uri: Optional[Union[URI, str]] = None,
@@ -68,9 +66,12 @@ class WebSocketProvider(PersistentConnectionProvider):
         # `PersistentConnectionProvider` kwargs can be passed through
         **kwargs: Any,
     ) -> None:
+        # initialize the endpoint_uri before calling the super constructor
         self.endpoint_uri = (
             URI(endpoint_uri) if endpoint_uri is not None else get_default_endpoint()
         )
+        super().__init__(**kwargs)
+        self._ws: Optional[WebSocketClientProtocol] = None
 
         if not any(
             self.endpoint_uri.startswith(prefix)
@@ -93,8 +94,6 @@ class WebSocketProvider(PersistentConnectionProvider):
 
         self.websocket_kwargs = merge(DEFAULT_WEBSOCKET_KWARGS, websocket_kwargs or {})
 
-        super().__init__(**kwargs)
-
     def __str__(self) -> str:
         return f"WebSocket connection: {self.endpoint_uri}"
 
@@ -133,6 +132,7 @@ class WebSocketProvider(PersistentConnectionProvider):
         self._ws = await connect(self.endpoint_uri, **self.websocket_kwargs)
 
     async def _provider_specific_disconnect(self) -> None:
+        # this should remain idempotent
         if self._ws is not None and not self._ws.closed:
             await self._ws.close()
             self._ws = None

web3/providers/rpc/async_rpc.py

@@ -168,12 +168,25 @@ class AsyncHTTPProvider(AsyncJSONBaseProvider):
 
     async def make_batch_request(
         self, batch_requests: List[Tuple[RPCEndpoint, Any]]
-    ) -> List[RPCResponse]:
+    ) -> Union[List[RPCResponse], RPCResponse]:
         self.logger.debug(f"Making batch request HTTP - uri: `{self.endpoint_uri}`")
         request_data = self.encode_batch_rpc_request(batch_requests)
         raw_response = await self._request_session_manager.async_make_post_request(
             self.endpoint_uri, request_data, **self.get_request_kwargs()
         )
         self.logger.debug("Received batch response HTTP.")
-        responses_list = cast(List[RPCResponse], self.decode_rpc_response(raw_response))
-        return sort_batch_response_by_response_ids(responses_list)
+        response = self.decode_rpc_response(raw_response)
+        if not isinstance(response, list):
+            # RPC errors return only one response with the error object
+            return response
+        return sort_batch_response_by_response_ids(
+            cast(List[RPCResponse], sort_batch_response_by_response_ids(response))
+        )
+
+    async def disconnect(self) -> None:
+        cache = self._request_session_manager.session_cache
+        for _, session in cache.items():
+            await session.close()
+        cache.clear()
+
+        self.logger.info(f"Successfully disconnected from: {self.endpoint_uri}")

web3/providers/rpc/rpc.py

@@ -71,6 +71,7 @@ class HTTPProvider(JSONBaseProvider):
         ] = empty,
         **kwargs: Any,
     ) -> None:
+        super().__init__(**kwargs)
         self._request_session_manager = HTTPSessionManager()
 
         if endpoint_uri is None:
@@ -88,8 +89,6 @@ class HTTPProvider(JSONBaseProvider):
                 self.endpoint_uri, session
             )
 
-        super().__init__(**kwargs)
-
     def __str__(self) -> str:
         return f"RPC connection {self.endpoint_uri}"
 
@@ -177,12 +176,17 @@ class HTTPProvider(JSONBaseProvider):
 
     def make_batch_request(
         self, batch_requests: List[Tuple[RPCEndpoint, Any]]
-    ) -> List[RPCResponse]:
+    ) -> Union[List[RPCResponse], RPCResponse]:
         self.logger.debug(f"Making batch request HTTP, uri: `{self.endpoint_uri}`")
         request_data = self.encode_batch_rpc_request(batch_requests)
         raw_response = self._request_session_manager.make_post_request(
             self.endpoint_uri, request_data, **self.get_request_kwargs()
         )
         self.logger.debug("Received batch response HTTP.")
-        responses_list = cast(List[RPCResponse], self.decode_rpc_response(raw_response))
-        return sort_batch_response_by_response_ids(responses_list)
+        response = self.decode_rpc_response(raw_response)
+        if not isinstance(response, list):
+            # RPC errors return only one response with the error object
+            return response
+        return sort_batch_response_by_response_ids(
+            cast(List[RPCResponse], sort_batch_response_by_response_ids(response))
+        )