Merge branch 'main' into gabrys/bundle-neptune-api

gabrys · web-flow · commit 0e056f506ba2 · 2026-01-13T12:46:00.000+01:00
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -1,30 +1,28 @@
 exclude: ^src/neptune_query/generated/
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.3.0
+    rev: v6.0.0
     hooks:
       - id: check-yaml
       - id: end-of-file-fixer
       - id: trailing-whitespace
   - repo: https://github.com/pycqa/isort
-    rev: 5.13.2
+    rev: 7.0.0
     hooks:
       - id: isort
         args: [--settings-path, pyproject.toml]
   - repo: https://github.com/psf/black
-    rev: 22.6.0
+    rev: 25.12.0
     hooks:
       - id: black
         args: [--config, pyproject.toml]
   - repo: https://github.com/pycqa/flake8
     rev: 7.0.0
     hooks:
       - id: flake8
-        entry: pflake8
-        # We use a custom version of flake8 that includes a fix for Python 3.12 compatibility
-        additional_dependencies: ["git+https://github.com/neptune-ai/pyproject-flake8.git@7.0.0"]
+        additional_dependencies: [Flake8-pyproject==1.2.4]
   - repo: https://github.com/Lucas-C/pre-commit-hooks
-    rev: v1.5.4
+    rev: v1.5.5
     hooks:
       - id: insert-license
         files: ^src/neptune_query.*[^/]+\.py$
diff --git a/src/neptune_query/filters.py b/src/neptune_query/filters.py
@@ -63,8 +63,7 @@ def any(*filters: "BaseAttributeFilter") -> "BaseAttributeFilter":
         return _AttributeFilterAlternative(filters=filters)
 
     @abc.abstractmethod
-    def _to_internal(self) -> _filters._BaseAttributeFilter:
-        ...
+    def _to_internal(self) -> _filters._BaseAttributeFilter: ...
 
 
 @dataclass
diff --git a/src/neptune_query/internal/composition/fetch_metric_buckets.py b/src/neptune_query/internal/composition/fetch_metric_buckets.py
@@ -64,8 +64,7 @@
 class _FetchInChunksProtocol(Protocol):
     def __call__(
         self, x_range: Optional[tuple[float, float]], bucket_limit: int
-    ) -> dict[RunAttributeDefinition, list[TimeseriesBucket]]:
-        ...
+    ) -> dict[RunAttributeDefinition, list[TimeseriesBucket]]: ...
 
 
 def fetch_metric_buckets(
@@ -252,9 +251,10 @@ def _compute_global_x_range(fetch_in_chunks: _FetchInChunksProtocol) -> Optional
     return x_range[0], x_range[1]
 
 
-def _update_range(
-    current_range: tuple[Optional[float], Optional[float]], bucket: TimeseriesBucket
-) -> tuple[Optional[float], Optional[float],]:
+def _update_range(current_range: tuple[Optional[float], Optional[float]], bucket: TimeseriesBucket) -> tuple[
+    Optional[float],
+    Optional[float],
+]:
     # We're including from_x and to_x because some buckets might hold only non-finite points,
     # in which case first_x and last_x are None.
     candidates = [bucket.first_x, bucket.last_x, bucket.from_x, bucket.to_x] + list(current_range)
diff --git a/src/neptune_query/internal/composition/fetch_metrics.py b/src/neptune_query/internal/composition/fetch_metrics.py
@@ -181,9 +181,9 @@ def go_fetch_sys_attrs() -> Generator[list[identifiers.SysId], None, None]:
         ),
     )
 
-    results: Generator[
-        dict[identifiers.RunAttributeDefinition, list[FloatPointValue]], None, None
-    ] = concurrency.gather_results(output)
+    results: Generator[dict[identifiers.RunAttributeDefinition, list[FloatPointValue]], None, None] = (
+        concurrency.gather_results(output)
+    )
 
     metrics_data: dict[identifiers.RunAttributeDefinition, list[FloatPointValue]] = {}
     for result in results:
diff --git a/src/neptune_query/internal/filters.py b/src/neptune_query/internal/filters.py
@@ -49,8 +49,7 @@ def any(filters: list["_BaseAttributeFilter"]) -> "_BaseAttributeFilter":
     @abc.abstractmethod
     def transform(
         self, map_attribute_filter: Callable[["_AttributeFilter"], "_AttributeFilter"]
-    ) -> "_BaseAttributeFilter":
-        ...
+    ) -> "_BaseAttributeFilter": ...
 
 
 @dataclass
@@ -351,8 +350,7 @@ def name_eq(name: str) -> "_Filter":
         return _Filter.eq(name_attribute, name)
 
     @abc.abstractmethod
-    def to_query(self) -> str:
-        ...
+    def to_query(self) -> str: ...
 
     def __str__(self) -> str:
         return self.to_query()
diff --git a/src/neptune_query/internal/retrieval/global_search.py b/src/neptune_query/internal/retrieval/global_search.py
@@ -117,13 +117,17 @@ def fetch_global_entries(
                     type=AttributeTypeDTO(
                         map_attribute_type_python_to_backend(sort_by.type) if sort_by.type is not None else "string"
                     ),
-                    aggregation_mode=QueryLeaderboardParamsFieldDTOAggregationMode(sort_by.aggregation)
-                    if sort_by.aggregation is not None
-                    else UNSET,
+                    aggregation_mode=(
+                        QueryLeaderboardParamsFieldDTOAggregationMode(sort_by.aggregation)
+                        if sort_by.aggregation is not None
+                        else UNSET
+                    ),
+                ),
+                dir_=(
+                    QueryLeaderboardParamsSortingParamsDTODir.ASCENDING
+                    if sort_direction == "asc"
+                    else QueryLeaderboardParamsSortingParamsDTODir.DESCENDING
                 ),
-                dir_=QueryLeaderboardParamsSortingParamsDTODir.ASCENDING
-                if sort_direction == "asc"
-                else QueryLeaderboardParamsSortingParamsDTODir.DESCENDING,
             ),
         ),
     )
diff --git a/src/neptune_query/internal/retrieval/metric_buckets.py b/src/neptune_query/internal/retrieval/metric_buckets.py
@@ -41,7 +41,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()
 
@@ -153,7 +153,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(
diff --git a/src/neptune_query/internal/retrieval/metrics.py b/src/neptune_query/internal/retrieval/metrics.py
@@ -39,7 +39,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,
diff --git a/src/neptune_query/internal/retrieval/search.py b/src/neptune_query/internal/retrieval/search.py
@@ -71,12 +71,10 @@ class ContainerType(Enum):
 
 class SysIdLabel(Protocol):
     @property
-    def sys_id(self) -> identifiers.SysId:
-        ...
+    def sys_id(self) -> identifiers.SysId: ...
 
     @property
-    def label(self) -> str:
-        ...
+    def label(self) -> str: ...
 
 
 @dataclass(frozen=True)
@@ -136,8 +134,7 @@ def __call__(
         limit: Optional[int] = None,
         batch_size: int = env.NEPTUNE_QUERY_SYS_ATTRS_BATCH_SIZE.get(),
         container_type: ContainerType = ContainerType.EXPERIMENT,
-    ) -> Generator[util.Page[T], None, None]:
-        ...
+    ) -> Generator[util.Page[T], None, None]: ...
 
 
 def _create_fetch_sys_attrs(
diff --git a/src/neptune_query/internal/retrieval/util.py b/src/neptune_query/internal/retrieval/util.py
@@ -57,13 +57,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()))
diff --git a/tests/fuzzy/test_fuzzy_output_format.py b/tests/fuzzy/test_fuzzy_output_format.py
@@ -61,9 +61,11 @@ def test_create_metrics_dataframe(
     # validate columns
     expected_attributes = sorted(
         {
-            f"{run_attributes.attribute_definition.name}:{run_attributes.attribute_definition.type}"
-            if type_suffix_in_column_names
-            else run_attributes.attribute_definition.name
+            (
+                f"{run_attributes.attribute_definition.name}:{run_attributes.attribute_definition.type}"
+                if type_suffix_in_column_names
+                else run_attributes.attribute_definition.name
+            )
             for run_attributes, points in metrics_data.items()
             if points
         }
diff --git a/tests/performance_e2e/backend/endpoints/get_multiple_float_series_values.py b/tests/performance_e2e/backend/endpoints/get_multiple_float_series_values.py
@@ -1,6 +1,7 @@
 """
 Endpoint for handling multiple float series values requests.
 """
+
 import hashlib
 import json
 import math
diff --git a/tests/performance_e2e/backend/endpoints/query_attribute_definitions_within_project.py b/tests/performance_e2e/backend/endpoints/query_attribute_definitions_within_project.py
@@ -1,6 +1,7 @@
 """
 Endpoint for handling attribute definitions queries within a project.
 """
+
 import json
 import random
 import time
diff --git a/tests/performance_e2e/backend/endpoints/search_leaderboard_entries.py b/tests/performance_e2e/backend/endpoints/search_leaderboard_entries.py
@@ -1,6 +1,7 @@
 """
 Endpoint for handling leaderboard entries search requests.
 """
+
 import json
 import random
 import time
diff --git a/tests/performance_e2e/backend/main.py b/tests/performance_e2e/backend/main.py
@@ -2,6 +2,7 @@
 Main entry point for the performance_e2e test backend server.
 Configures the FastAPI app and includes all endpoints.
 """
+
 import time
 from typing import Dict
 
diff --git a/tests/performance_e2e/backend/middleware/add_latency_middleware.py b/tests/performance_e2e/backend/middleware/add_latency_middleware.py
@@ -1,6 +1,7 @@
 """
 Middleware for simulating network latency in performance_e2e tests.
 """
+
 import asyncio
 import random
 import time
diff --git a/tests/performance_e2e/backend/middleware/read_perf_config_middleware.py b/tests/performance_e2e/backend/middleware/read_perf_config_middleware.py
@@ -2,6 +2,7 @@
 Middleware for handling performance_e2e testing configuration parsing.
 Parses X-Perf-Request headers and attaches config to request state.
 """
+
 import json
 from typing import (
     Callable,
diff --git a/tests/performance_e2e/backend/middleware/request_logging_middleware.py b/tests/performance_e2e/backend/middleware/request_logging_middleware.py
@@ -2,6 +2,7 @@
 Middleware for tracking request metrics in performance_e2e tests.
 Records timing information and adds metrics to requests.
 """
+
 import json
 import time
 from contextvars import Token
diff --git a/tests/performance_e2e/backend/perf_request.py b/tests/performance_e2e/backend/perf_request.py
@@ -2,6 +2,7 @@
 Common schema and utilities for X-Perf-Request headers used in performance_e2e testing.
 This module provides shared functionality for both test cases and the server implementation.
 """
+
 import json
 from dataclasses import (
     asdict,
diff --git a/tests/performance_e2e/backend/utils/logging.py b/tests/performance_e2e/backend/utils/logging.py
@@ -2,6 +2,7 @@
 Logging utilities for the performance_e2e test backend.
 Provides a consistent logging configuration across all modules.
 """
+
 import logging
 import os
 from contextvars import ContextVar
diff --git a/tests/performance_e2e/backend/utils/metrics.py b/tests/performance_e2e/backend/utils/metrics.py
@@ -1,6 +1,7 @@
 """
 Metrics utilities for performance_e2e monitoring.
 """
+
 from dataclasses import dataclass
 
 
diff --git a/tests/performance_e2e/backend/utils/random_utils.py b/tests/performance_e2e/backend/utils/random_utils.py
@@ -1,6 +1,7 @@
 """
 Utilities for generating random data for test responses.
 """
+
 import random
 import string
 from typing import Final
diff --git a/tests/unit/internal/retrieval/test_protobuf_payload.py b/tests/unit/internal/retrieval/test_protobuf_payload.py
@@ -8,14 +8,14 @@
     ProtoView,
 )
 from neptune_query.generated.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),
@@ -34,5 +34,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

Original file line number	Diff line number	Diff line change
`@@ -39,7 +39,7 @@`
`39`	`39`	`logger = get_logger()`
`40`	`40`
`41`	`41`	`# Tuples are used here to enhance performance`
`42`		`-FloatPointValue = tuple[float, float, float, bool, float]`
	`42`	`+FloatPointValue = tuple[int, float, float, bool, float]`
`43`	`43`	`(`
`44`	`44`	`TimestampIndex,`
`45`	`45`	`StepIndex,`