Fix arcane mypy errors (#189)

Author: gabrys

src/neptune_query/internal/retrieval/metric_buckets.py

@@ -39,7 +39,7 @@
 from ..query_metadata_context import with_neptune_client_metadata
 from . import retry
 from .search import ContainerType
-from .util import ProtobufPayload
+from .util import body_from_protobuf
 
 logger = get_logger()
 
@@ -151,7 +151,7 @@ def fetch_time_series_buckets(
     call_api = retry.handle_errors_default(with_neptune_client_metadata(get_timeseries_buckets_proto.sync_detailed))
     response = call_api(
         client=client,
-        body=ProtobufPayload(request_object),
+        body=body_from_protobuf(request_object),
     )
 
     logger.debug(

src/neptune_query/internal/retrieval/metrics.py

@@ -38,7 +38,7 @@
 logger = get_logger()
 
 # Tuples are used here to enhance performance
-FloatPointValue = tuple[float, float, float, bool, float]
+FloatPointValue = tuple[int, float, float, bool, float]
 (
     TimestampIndex,
     StepIndex,

src/neptune_query/internal/retrieval/util.py

@@ -56,13 +56,19 @@ def fetch_pages(
         page_params = make_new_page_params(page_params, data)
 
 
-class ProtobufPayload(File):
-    """A version of the neptune_api.types.File class that uses a protobuf message as payload"""
+class ReusableFile(File):
+    """A File that recreates its payload on each access to support retries."""
 
     @property
     def payload(self) -> BinaryIO:
         return self.get_payload()
 
     @payload.setter
-    def payload(self, message: Message) -> None:
-        self.get_payload = lambda: BytesIO(message.SerializeToString())
+    def payload(self, payload: BinaryIO) -> None:
+        stored_payload = payload.read()
+        self.get_payload = lambda: BytesIO(stored_payload)
+
+
+def body_from_protobuf(message: Message) -> ReusableFile:
+    """Generate a ReusableFile from a protobuf message."""
+    return ReusableFile(payload=BytesIO(message.SerializeToString()))

tests/unit/internal/retrieval/test_protobuf_payload.py

@@ -9,14 +9,14 @@
 )
 from neptune_api.types import File
 
-from neptune_query.internal.retrieval.util import ProtobufPayload
+from neptune_query.internal.retrieval.util import body_from_protobuf
 
 
 @mock.patch("neptune_query.internal.retrieval.util.BytesIO", wraps=BytesIO)
 def test_bytesio_recreation_on_retry(mock_bytesio):
     """Test that BytesIO instance is recreated when the API call is retried."""
 
-    file = ProtobufPayload(
+    file = body_from_protobuf(
         ProtoGetTimeseriesBucketsRequest(
             expressions=[ProtoCustomExpression(requestId="0123", customYFormula="${abc}")],
             view=ProtoView(xScale=ProtoScale.linear),
@@ -35,5 +35,6 @@ def test_bytesio_recreation_on_retry(mock_bytesio):
     assert read1 == read2
     assert read2 != b""
 
-    # Verify BytesIO was called twice
-    assert mock_bytesio.call_count == 2
+    # Verify BytesIO was called three times:
+    # Once during ProtobufPayload creation and twice for the two reads
+    assert mock_bytesio.call_count == 3