aiohttp_rpc.client.websocket.WsJsonRpcClient._handle_json_responses() - Code Metrics - expert-m/aiohttp-rpc - Measure and Improve Code Quality continuously with Scrutinizer

WsJsonRpcClient._handle_json_responses() B
last analyzed 2023-10-15 08:14 UTC

↳ Parent: aiohttp_rpc.client.websocket

Complexity

Conditions

Size

Total Lines	14
Code Lines	12

Duplication

Lines	0
Ratio	0 %

Importance

Changes

Metric	Value
eloc	12
dl	0
loc	14
rs	8.6666
c	0
b	0
f	0
cc	6
nop	4

import asyncio
import logging
import typing

from aiohttp import ClientSession, http_websocket, web_ws

from .base import BaseJsonRpcClient
from .. import errors, typedefs, utils


__all__ = (
    'WsJsonRpcClient',
)

logger = logging.getLogger(__name__)


class WsJsonRpcClient(BaseJsonRpcClient):
    url: typing.Optional[str]
    ws_connect: typing.Optional[typedefs.WSConnectType]
    ws_connect_kwargs: dict
    _timeout: typing.Optional[float]
    _timeout_for_data_receiving: typing.Optional[float]
    _connection_check_interval: typing.Optional[float]
    _pending: typing.Dict[typing.Any, asyncio.Future]
    _message_worker: typing.Optional[asyncio.Future] = None
    _check_worker: typing.Optional[asyncio.Future] = None
    _session_is_outer: bool
    _ws_connect_is_outer: bool
    _json_request_handler: typing.Optional[typing.Callable] = None
    _unprocessed_json_response_handler: typing.Optional[typing.Callable] = None
    _background_tasks: typing.Set
    _is_closed: bool = True

    def __init__(self,
                 url: typing.Optional[str] = None, *,
                 session: typing.Optional[ClientSession] = None,
                 ws_connect: typing.Optional[typedefs.WSConnectType] = None,
                 timeout: typing.Optional[float] = 60,
                 timeout_for_data_receiving: typing.Optional[float] = None,
                 connection_check_interval: typing.Optional[float] = 5,
                 json_request_handler: typing.Optional[typing.Callable] = None,
                 unprocessed_json_response_handler: typing.Optional[typing.Callable] = None,
                 **ws_connect_kwargs) -> None:
        assert (session is not None) or (url is not None and session is None) or (ws_connect is not None)

        self.url = url
        self._timeout = timeout
        self._timeout_for_data_receiving = timeout_for_data_receiving
        self._connection_check_interval = connection_check_interval

        self.session = session
        self._session_is_outer = session is not None  # We don't close an outer session.

        self.ws_connect = ws_connect
        self.ws_connect_kwargs = ws_connect_kwargs
        self._ws_connect_is_outer = ws_connect is not None  # We don't close an outer WS connection.

        self._pending = {}
        self._json_request_handler = json_request_handler
        self._unprocessed_json_response_handler = unprocessed_json_response_handler
        self._background_tasks = set()

    async def connect(self) -> None:
        self._is_closed = False

        if self.session is None and self.ws_connect is None:
            self.session = ClientSession(json_serialize=self.json_serialize)

        if self.ws_connect is None:
            assert self.url is not None and self.session is not None

            try:
                self.ws_connect = await self.session.ws_connect(self.url, **self.ws_connect_kwargs)
            except Exception:
                await self.disconnect()
                raise

        self._message_worker = asyncio.create_task(self._handle_ws_messages())

        if self._connection_check_interval is not None:
            self._check_worker = asyncio.create_task(self._check_ws_connection())

    async def disconnect(self) -> None:
        self._is_closed = True

        if self.ws_connect is not None and not self._ws_connect_is_outer:
            await self.ws_connect.close()

        if self.session is not None and not self._session_is_outer:
            await self.session.close()

        if self._message_worker is not None:
            if self._ws_connect_is_outer:
                await asyncio.wait_for(self._message_worker, timeout=60)
            else:
                await self._message_worker

        if self._check_worker is not None:
            self._check_worker.cancel()
            await self._check_worker

    async def send_json(self,
                        data: typing.Any, *,
                        without_response: bool = False,
                        **kwargs) -> typing.Tuple[typing.Any, typing.Optional[dict]]:
        assert self.ws_connect is not None

        if without_response:
            try:
                await self.ws_connect.send_str(self.json_serialize(data), **kwargs)
            except ConnectionResetError as e:
                error = errors.ServerError(utils.get_exc_message(e)).with_traceback()
                self._notify_all_about_error(error)
                raise error

            return None, None

        request_ids = self._get_ids_from_json(data)
        future: asyncio.Future = asyncio.Future()

        for request_id in request_ids:
            self._pending[request_id] = future

        try:
            await self.ws_connect.send_str(self.json_serialize(data), **kwargs)
        except ConnectionResetError as e:
            error = errors.ServerError(utils.get_exc_message(e)).with_traceback()
            self._notify_all_about_error(error)
            raise error

        if not request_ids:
            return None, None

        if self._timeout is not None:
            future = asyncio.wait_for(future, timeout=self._timeout)  # type: ignore

        result = await future

        return result, None

    @staticmethod
    def _get_ids_from_json(data: typing.Any) -> typing.Tuple[typedefs.JsonRpcIdType, ...]:
        if not data:
            return ()

        if isinstance(data, typing.Mapping) and data.get('id') is not None:
            return (
                data['id'],
            )

        if isinstance(data, typing.Sequence):
            return tuple(
                item['id']
                for item in data
                if isinstance(item, typing.Mapping) and item.get('id') is not None
            )

        return ()

    async def _handle_ws_messages(self) -> None:
        assert self.ws_connect is not None

        while True:
            try:
                ws_msg: http_websocket.WSMessage = await self.ws_connect.receive(
                    timeout=self._timeout_for_data_receiving,
                )
            except asyncio.TimeoutError:
                if self._is_closed:
                    break
                else:
                    continue

            if ws_msg.type in (
                http_websocket.WSMsgType.CLOSE,
                http_websocket.WSMsgType.CLOSING,
                http_websocket.WSMsgType.CLOSED,
            ):
                break

            if ws_msg.type != http_websocket.WSMsgType.TEXT:
                continue

            try:
                task = asyncio.create_task(self._handle_single_ws_message(ws_msg))
            except asyncio.CancelledError as e:
                error = errors.InternalError(utils.get_exc_message(e)).with_traceback()
                self._notify_all_about_error(error)
                break
            except Exception:
                logger.warning('Can\'t process WS message.', exc_info=True)
            else:
                # To avoid a task disappearing mid execution:
                self._background_tasks.add(task)
                task.add_done_callback(self._background_tasks.discard)

            if self._is_closed:
                break

    async def _check_ws_connection(self) -> None:
        assert self.ws_connect is not None

        try:
            while not self._is_closed:
                if self.ws_connect.closed:
                    error = errors.ServerError('Connection is closed')
                    self._notify_all_about_error(error)
                    break

                await asyncio.sleep(self._connection_check_interval)  # type: ignore
        except asyncio.CancelledError:
            pass

    async def _handle_single_ws_message(self, ws_msg: http_websocket.WSMessage) -> None:
        if ws_msg.type != http_websocket.WSMsgType.text:
            return

        try:
            json_response = self.json_deserialize(ws_msg.data)
        except Exception:
            logger.warning('Can\'t parse json.', exc_info=True)
            return

        if not json_response:
            return

        if isinstance(json_response, typing.Mapping):
            await self._handle_single_json_response(json_response, ws_msg=ws_msg)
            return

        if isinstance(json_response, typing.Sequence):
            await self._handle_json_responses(json_response, ws_msg=ws_msg)
            return

        logger.warning('Couldn\'t process the response.', extra={
            'json_response': json_response,
        })

    async def _handle_single_json_response(self, json_response: typing.Mapping, *, ws_msg: web_ws.WSMessage) -> None:
        if 'method' in json_response:
            if self._json_request_handler is not None:
                await self._json_request_handler(
                    ws_connect=self.ws_connect,
                    ws_msg=ws_msg,
                    json_request=json_response,
                )
        elif 'id' in json_response and json_response['id'] in self._pending:
            self._notify_about_result(json_response['id'], json_response)
        elif self._unprocessed_json_response_handler is not None:
            self._unprocessed_json_response_handler(
                ws_connect=self.ws_connect,
                ws_msg=ws_msg,
                json_response=json_response,
            )

    async def _handle_json_responses(self, json_responses: typing.Sequence, *, ws_msg: web_ws.WSMessage) -> None:
        if isinstance(json_responses[0], typing.Mapping) and 'method' in json_responses[0]:
            if self._json_request_handler is not None:
                await self._json_request_handler(ws_connect=self.ws_connect, ws_msg=ws_msg)
        else:
            response_ids = self._get_ids_from_json(json_responses)

            if response_ids:
                self._notify_about_results(response_ids, json_responses)
            elif self._unprocessed_json_response_handler is not None:
                self._unprocessed_json_response_handler(
                    ws_connect=self.ws_connect,
                    ws_msg=ws_msg,
                    json_response=json_responses,
                )

    def _notify_all_about_error(self, error: Exception) -> None:
        for future in self._pending.values():
            try:
                future.set_exception(error)
            except asyncio.InvalidStateError:
                pass

        self._pending.clear()

    def _notify_about_result(self, response_id: typedefs.JsonRpcIdType, json_response: typing.Mapping) -> None:
        future = self._pending.pop(response_id, None)

        if future is not None:
            future.set_result(json_response)

    def _notify_about_results(self,
                              response_ids: typing.Sequence[typedefs.JsonRpcIdType],
                              json_response: typing.Sequence) -> None:
        is_processed = False

        for response_id in response_ids:
            future = self._pending.pop(response_id, None)

            if future is not None and not is_processed:
                # We suppose that a batch result has the same ids that we sent.
                # And these ids have the same future.

                future.set_result(json_response)
                is_processed = True


1			import asyncio
2			import logging
3			import typing
4
5			from aiohttp import ClientSession, http_websocket, web_ws
6
7			from .base import BaseJsonRpcClient
8			from .. import errors, typedefs, utils
9
10
11			__all__ = (
12			'WsJsonRpcClient',
13			)
14
15			logger = logging.getLogger(__name__)
16
17
18			class WsJsonRpcClient(BaseJsonRpcClient):
19			url: typing.Optional[str]
20			ws_connect: typing.Optional[typedefs.WSConnectType]
21			ws_connect_kwargs: dict
22			_timeout: typing.Optional[float]
23			_timeout_for_data_receiving: typing.Optional[float]
24			_connection_check_interval: typing.Optional[float]
25			_pending: typing.Dict[typing.Any, asyncio.Future]
26			_message_worker: typing.Optional[asyncio.Future] = None
27			_check_worker: typing.Optional[asyncio.Future] = None
28			_session_is_outer: bool
29			_ws_connect_is_outer: bool
30			_json_request_handler: typing.Optional[typing.Callable] = None
31			_unprocessed_json_response_handler: typing.Optional[typing.Callable] = None
32			_background_tasks: typing.Set
33			_is_closed: bool = True
34
35			def __init__(self,
36			url: typing.Optional[str] = None, *,
37			session: typing.Optional[ClientSession] = None,
38			ws_connect: typing.Optional[typedefs.WSConnectType] = None,
39			timeout: typing.Optional[float] = 60,
40			timeout_for_data_receiving: typing.Optional[float] = None,
41			connection_check_interval: typing.Optional[float] = 5,
42			json_request_handler: typing.Optional[typing.Callable] = None,
43			unprocessed_json_response_handler: typing.Optional[typing.Callable] = None,
44			**ws_connect_kwargs) -> None:
45			assert (session is not None) or (url is not None and session is None) or (ws_connect is not None)
46
47			self.url = url
48			self._timeout = timeout
49			self._timeout_for_data_receiving = timeout_for_data_receiving
50			self._connection_check_interval = connection_check_interval
51
52			self.session = session
53			self._session_is_outer = session is not None # We don't close an outer session.
54
55			self.ws_connect = ws_connect
56			self.ws_connect_kwargs = ws_connect_kwargs
57			self._ws_connect_is_outer = ws_connect is not None # We don't close an outer WS connection.
58
59			self._pending = {}
60			self._json_request_handler = json_request_handler
61			self._unprocessed_json_response_handler = unprocessed_json_response_handler
62			self._background_tasks = set()
63
64			async def connect(self) -> None:
65			self._is_closed = False
66
67			if self.session is None and self.ws_connect is None:
68			self.session = ClientSession(json_serialize=self.json_serialize)
69
70			if self.ws_connect is None:
71			assert self.url is not None and self.session is not None
72
73			try:
74			self.ws_connect = await self.session.ws_connect(self.url, **self.ws_connect_kwargs)
75			except Exception:
76			await self.disconnect()
77			raise
78
79			self._message_worker = asyncio.create_task(self._handle_ws_messages())
80
81			if self._connection_check_interval is not None:
82			self._check_worker = asyncio.create_task(self._check_ws_connection())
83
84			async def disconnect(self) -> None:
85			self._is_closed = True
86
87			if self.ws_connect is not None and not self._ws_connect_is_outer:
88			await self.ws_connect.close()
89
90			if self.session is not None and not self._session_is_outer:
91			await self.session.close()
92
93			if self._message_worker is not None:
94			if self._ws_connect_is_outer:
95			await asyncio.wait_for(self._message_worker, timeout=60)
96			else:
97			await self._message_worker
98
99			if self._check_worker is not None:
100			self._check_worker.cancel()
101			await self._check_worker
102
103			async def send_json(self,
104			data: typing.Any, *,
105			without_response: bool = False,
106			**kwargs) -> typing.Tuple[typing.Any, typing.Optional[dict]]:
107			assert self.ws_connect is not None
108
109			if without_response:
110			try:
111			await self.ws_connect.send_str(self.json_serialize(data), **kwargs)
112			except ConnectionResetError as e:
113			error = errors.ServerError(utils.get_exc_message(e)).with_traceback()
114			self._notify_all_about_error(error)
115			raise error
116
117			return None, None
118
119			request_ids = self._get_ids_from_json(data)
120			future: asyncio.Future = asyncio.Future()
121
122			for request_id in request_ids:
123			self._pending[request_id] = future
124
125			try:
126			await self.ws_connect.send_str(self.json_serialize(data), **kwargs)
127			except ConnectionResetError as e:
128			error = errors.ServerError(utils.get_exc_message(e)).with_traceback()
129			self._notify_all_about_error(error)
130			raise error
131
132			if not request_ids:
133			return None, None
134
135			if self._timeout is not None:
136			future = asyncio.wait_for(future, timeout=self._timeout) # type: ignore
137
138			result = await future
139
140			return result, None
141
142			@staticmethod
143			def _get_ids_from_json(data: typing.Any) -> typing.Tuple[typedefs.JsonRpcIdType, ...]:
144			if not data:
145			return ()
146
147			if isinstance(data, typing.Mapping) and data.get('id') is not None:
148			return (
149			data['id'],
150			)
151
152			if isinstance(data, typing.Sequence):
153			return tuple(
154			item['id']
155			for item in data
156			if isinstance(item, typing.Mapping) and item.get('id') is not None
157			)
158
159			return ()
160
161			async def _handle_ws_messages(self) -> None:
162			assert self.ws_connect is not None
163
164			while True:
165			try:
166			ws_msg: http_websocket.WSMessage = await self.ws_connect.receive(
167			timeout=self._timeout_for_data_receiving,
168			)
169			except asyncio.TimeoutError:
170			if self._is_closed:
171			break
172			else:
173			continue
174
175			if ws_msg.type in (
176			http_websocket.WSMsgType.CLOSE,
177			http_websocket.WSMsgType.CLOSING,
178			http_websocket.WSMsgType.CLOSED,
179			):
180			break
181
182			if ws_msg.type != http_websocket.WSMsgType.TEXT:
183			continue
184
185			try:
186			task = asyncio.create_task(self._handle_single_ws_message(ws_msg))
187			except asyncio.CancelledError as e:
188			error = errors.InternalError(utils.get_exc_message(e)).with_traceback()
189			self._notify_all_about_error(error)
190			break
191			except Exception:
192			logger.warning('Can\'t process WS message.', exc_info=True)
193			else:
194			# To avoid a task disappearing mid execution:
195			self._background_tasks.add(task)
196			task.add_done_callback(self._background_tasks.discard)
197
198			if self._is_closed:
199			break
200
201			async def _check_ws_connection(self) -> None:
202			assert self.ws_connect is not None
203
204			try:
205			while not self._is_closed:
206			if self.ws_connect.closed:
207			error = errors.ServerError('Connection is closed')
208			self._notify_all_about_error(error)
209			break
210
211			await asyncio.sleep(self._connection_check_interval) # type: ignore
212			except asyncio.CancelledError:
213			pass
214
215			async def _handle_single_ws_message(self, ws_msg: http_websocket.WSMessage) -> None:
216			if ws_msg.type != http_websocket.WSMsgType.text:
217			return
218
219			try:
220			json_response = self.json_deserialize(ws_msg.data)
221			except Exception:
222			logger.warning('Can\'t parse json.', exc_info=True)
223			return
224
225			if not json_response:
226			return
227
228			if isinstance(json_response, typing.Mapping):
229			await self._handle_single_json_response(json_response, ws_msg=ws_msg)
230			return
231
232			if isinstance(json_response, typing.Sequence):
233			await self._handle_json_responses(json_response, ws_msg=ws_msg)
234			return
235
236			logger.warning('Couldn\'t process the response.', extra={
237			'json_response': json_response,
238			})
239
240			async def _handle_single_json_response(self, json_response: typing.Mapping, *, ws_msg: web_ws.WSMessage) -> None:
241			if 'method' in json_response:
242			if self._json_request_handler is not None:
243			await self._json_request_handler(
244			ws_connect=self.ws_connect,
245			ws_msg=ws_msg,
246			json_request=json_response,
247			)
248			elif 'id' in json_response and json_response['id'] in self._pending:
249			self._notify_about_result(json_response['id'], json_response)
250			elif self._unprocessed_json_response_handler is not None:
251			self._unprocessed_json_response_handler(
252			ws_connect=self.ws_connect,
253			ws_msg=ws_msg,
254			json_response=json_response,
255			)
256
257			async def _handle_json_responses(self, json_responses: typing.Sequence, *, ws_msg: web_ws.WSMessage) -> None:
258			if isinstance(json_responses[0], typing.Mapping) and 'method' in json_responses[0]:
259			if self._json_request_handler is not None:
260			await self._json_request_handler(ws_connect=self.ws_connect, ws_msg=ws_msg)
261			else:
262			response_ids = self._get_ids_from_json(json_responses)
263
264			if response_ids:
265			self._notify_about_results(response_ids, json_responses)
266			elif self._unprocessed_json_response_handler is not None:
267			self._unprocessed_json_response_handler(
268			ws_connect=self.ws_connect,
269			ws_msg=ws_msg,
270			json_response=json_responses,
271			)
272
273			def _notify_all_about_error(self, error: Exception) -> None:
274			for future in self._pending.values():
275			try:
276			future.set_exception(error)
277			except asyncio.InvalidStateError:
278			pass
279
280			self._pending.clear()
281
282			def _notify_about_result(self, response_id: typedefs.JsonRpcIdType, json_response: typing.Mapping) -> None:
283			future = self._pending.pop(response_id, None)
284
285			if future is not None:
286			future.set_result(json_response)
287
288			def _notify_about_results(self,
289			response_ids: typing.Sequence[typedefs.JsonRpcIdType],
290			json_response: typing.Sequence) -> None:
291			is_processed = False
292
293			for response_id in response_ids:
294			future = self._pending.pop(response_id, None)
295
296			if future is not None and not is_processed:
297			# We suppose that a batch result has the same ids that we sent.
298			# And these ids have the same future.
299
300			future.set_result(json_response)
301			is_processed = True
302

expert-m / aiohttp-rpc

WsJsonRpcClient._handle_json_responses() B last analyzed 2023-10-15 08:14 UTC

Complexity

Size

Duplication

Importance

Duplication Side-by-Side

Filter issues like

WsJsonRpcClient._handle_json_responses() B
last analyzed 2023-10-15 08:14 UTC