review fixes

Author: vgvoleg

ydb/_topic_reader/datatypes.py

@@ -7,7 +7,7 @@
 from collections import deque
 from dataclasses import dataclass, field
 import datetime
-from typing import Union, Any, List, Dict, Deque, Optional
+from typing import Union, Any, List, Dict, Deque, Optional, Tuple
 
 from ydb._grpc.grpcwrapper.ydb_topic import OffsetsRange, Codec
 from ydb._topic_reader import topic_reader_asyncio
@@ -171,3 +171,11 @@ def alive(self) -> bool:
 
     def pop_message(self) -> PublicMessage:
         return self.messages.pop(0)
+
+    def _extend(self, batch: PublicBatch) -> None:
+        self.messages.extend(batch.messages)
+        self._bytes_size += batch._bytes_size
+
+    def _pop(self) -> Tuple[List[PublicMessage], bool]:
+        msgs_left = True if len(self.messages) > 1 else False
+        return self.messages.pop(0), msgs_left

ydb/_topic_reader/topic_reader_asyncio.py

@@ -384,9 +384,9 @@ def receive_message_nowait(self):
 
         part_sess_id, batch = self._get_first_batch()
 
-        message = batch.messages.pop(0)
+        message, msgs_left = batch._pop()
 
-        if len(batch.messages) == 0:
+        if not msgs_left:
             self._buffer_release_bytes(batch._bytes_size)
         else:
             # TODO: we should somehow release bytes from single message as well
@@ -620,8 +620,7 @@ async def _decode_batches_loop(self):
     def _add_batch_to_queue(self, batch: datatypes.PublicBatch):
         part_sess_id = batch._partition_session.id
         if part_sess_id in self._message_batches:
-            self._message_batches[part_sess_id].messages.extend(batch.messages)
-            self._message_batches[part_sess_id]._bytes_size += batch._bytes_size
+            self._message_batches[part_sess_id]._extend(batch)
             return
 
         self._message_batches[part_sess_id] = batch