Specify types in set_parser() (#12073)

Author: Dreamsorcerer

aiohttp/client_proto.py

@@ -1,7 +1,8 @@
 import asyncio
 from contextlib import suppress
-from typing import Any, Callable
+from typing import Callable, Protocol
 
+from ._websocket.reader import WebSocketDataQueue
 from .base_protocol import BaseProtocol
 from .client_exceptions import (
     ClientConnectionError,
@@ -14,6 +15,7 @@
     _EXC_SENTINEL,
     EMPTY_BODY_STATUS_CODES,
     BaseTimerContext,
+    ErrorableProtocol,
     set_exception,
     set_result,
 )
@@ -22,6 +24,10 @@
 from .streams import EMPTY_PAYLOAD, DataQueue, StreamReader
 
 
+class _Payload(ErrorableProtocol, Protocol):
+    def is_eof(self) -> bool: ...
+
+
 class ResponseHandler(BaseProtocol, DataQueue[tuple[RawResponseMessage, StreamReader]]):
     """Helper class to adapt between Protocol and StreamReader."""
 
@@ -31,7 +37,7 @@ def __init__(self, loop: asyncio.AbstractEventLoop) -> None:
 
         self._should_close = False
 
-        self._payload: StreamReader | None = None
+        self._payload: _Payload | None = None
         self._skip_payload = False
         self._payload_parser: WebSocketReader | None = None
         self._data_received_cb: Callable[[], None] | None = None
@@ -206,15 +212,10 @@ def set_exception(
 
     def set_parser(
         self,
-        parser: Any,
-        payload: Any,
+        parser: WebSocketReader,
+        payload: WebSocketDataQueue,
         data_received_cb: Callable[[], None] | None = None,
     ) -> None:
-        # TODO: actual types are:
-        #   parser: WebSocketReader
-        #   payload: WebSocketDataQueue
-        # but they are not generi enough
-        # Need an ABC for both types
         self._payload = payload
         self._payload_parser = parser
         self._data_received_cb = data_received_cb