pre-commit fixes

Author: klpoland

gateway/sds_gateway/api_methods/helpers/temporal_filtering.py

@@ -1,12 +1,15 @@
 import re
 
 from django.db.models import QuerySet
-
+from loguru import logger as log
 from opensearchpy.exceptions import NotFoundError as OpenSearchNotFoundError
-from sds_gateway.api_methods.models import CaptureType, Capture, File, DRF_RF_FILENAME_REGEX_STR
+
+from sds_gateway.api_methods.models import DRF_RF_FILENAME_REGEX_STR
+from sds_gateway.api_methods.models import Capture
+from sds_gateway.api_methods.models import CaptureType
+from sds_gateway.api_methods.models import File
 from sds_gateway.api_methods.utils.opensearch_client import get_opensearch_client
 from sds_gateway.api_methods.utils.relationship_utils import get_capture_files
-from loguru import logger as log
 
 # Digital RF spec: rf@SECONDS.MILLISECONDS.h5 (e.g. rf@1396379502.000.h5)
 # https://github.com/MITHaystack/digital_rf
@@ -55,14 +58,12 @@ def get_capture_bounds(capture_type: CaptureType, capture_uuid: str) -> tuple[in
     try:
         response = client.get(index=index, id=capture_uuid)
     except OpenSearchNotFoundError as e:
-        raise ValueError(
-            f"Capture {capture_uuid} not found in OpenSearch index {index}"
-        ) from e
+        msg = f"Capture {capture_uuid} not found in OpenSearch index {index}"
+        raise ValueError(msg) from e
 
     if not response.get("found"):
-        raise ValueError(
-            f"Capture {capture_uuid} not found in OpenSearch index {index}"
-        )
+        msg = f"Capture {capture_uuid} not found in OpenSearch index {index}"
+        raise ValueError(msg)
 
     source = response.get("_source", {})
     search_props = source.get("search_props", {})
@@ -90,7 +91,7 @@ def filter_capture_data_files_selection_bounds(
     capture_type: CaptureType,
     capture: Capture,
     start_time: int,  # relative ms from start of capture (from UI)
-    end_time: int,    # relative ms from start of capture (from UI)
+    end_time: int,  # relative ms from start of capture (from UI)
 ) -> QuerySet[File]:
     """Filter the capture file selection bounds to the given start and end times."""
     _catch_capture_type_error(capture_type)
@@ -108,26 +109,32 @@ def filter_capture_data_files_selection_bounds(
         name__lte=end_file_name,
     ).order_by("name")
 
+
 def get_capture_files_with_temporal_filter(
     capture_type: CaptureType,
     capture: Capture,
-    start_time: int | None = None, # milliseconds since start of capture
+    start_time: int | None = None,  # milliseconds since start of capture
     end_time: int | None = None,
 ) -> QuerySet[File]:
     """Get the capture files with temporal filtering."""
     _catch_capture_type_error(capture_type)
 
     if start_time is None or end_time is None:
-        log.warning("Start or end time is None, returning all capture files without temporal filtering")
+        log.warning(
+            "Start or end time is None; returning all capture files without "
+            "temporal filtering"
+        )
         return get_capture_files(capture)
 
     # get non-data files
-    non_data_files = get_capture_files(capture).exclude(name__regex=DRF_RF_FILENAME_REGEX_STR)
+    non_data_files = get_capture_files(capture).exclude(
+        name__regex=DRF_RF_FILENAME_REGEX_STR
+    )
 
     # get data files with temporal filtering
     data_files = filter_capture_data_files_selection_bounds(
         capture_type, capture, start_time, end_time
     )
 
     # return all files
-    return non_data_files.union(data_files)
+    return non_data_files.union(data_files)

gateway/sds_gateway/api_methods/models.py

@@ -13,10 +13,10 @@
 from blake3 import blake3 as Blake3  # noqa: N812
 from django.conf import settings
 from django.db import models
-from django.db.models import Sum
 from django.db.models import Count
 from django.db.models import ProtectedError
 from django.db.models import QuerySet
+from django.db.models import Sum
 from django.db.models.signals import post_save
 from django.db.models.signals import pre_delete
 from django.dispatch import receiver
@@ -423,22 +423,27 @@ def get_capture(self) -> dict[str, Any]:
             "owner": self.owner,
         }
 
-
     def get_drf_data_files_queryset(self) -> QuerySet[File]:
         """DRF data files (rf@*.h5) for this capture (M2M + FK)."""
         if self.capture_type != CaptureType.DigitalRF:
             log.warning("Capture %s is not a DigitalRF capture", self.uuid)
             return File.objects.none()
 
         # Local import avoids circular import (relationship_utils imports Capture).
-        from sds_gateway.api_methods.utils.relationship_utils import get_capture_files
+        from sds_gateway.api_methods.utils.relationship_utils import (  # noqa: PLC0415
+            get_capture_files,
+        )
 
         return get_capture_files(self, include_deleted=False).filter(
             name__regex=DRF_RF_FILENAME_REGEX_STR,
         )
 
     def get_drf_data_files_stats(self) -> dict[str, int]:
-        """Count + total size in one query; cached per instance. File PK is ``uuid`` — use ``pk``."""
+        """
+        Count + total size in one query; cached per instance.
+
+        File primary key is ``uuid``; use ``pk`` in aggregates.
+        """
         if hasattr(self, "_drf_data_files_stats_cache"):
             return self._drf_data_files_stats_cache
 

gateway/sds_gateway/api_methods/serializers/capture_serializers.py

@@ -99,7 +99,7 @@ def get_files(self, capture: Capture) -> ReturnList[File]:
     @extend_schema_field(serializers.IntegerField(allow_null=True))
     def get_total_file_size(self, capture: Capture) -> int | None:
         """Get the total file size of all files associated with this capture."""
-        
+
         if capture.capture_type != CaptureType.DigitalRF:
             return None
 
@@ -109,11 +109,16 @@ def get_total_file_size(self, capture: Capture) -> int | None:
         data_total = self.get_data_files_info(capture).get("total_size", 0)
         if total < data_total:
             logging.getLogger(__name__).warning(
-                "Capture %s: total_file_size (%s) < data_files_total_size (%s); using data total.",
-                str(capture.uuid), total, data_total,
+                (
+                    "Capture %s: total_file_size (%s) < data_files_total_size (%s); "
+                    "using data total."
+                ),
+                str(capture.uuid),
+                total,
+                data_total,
             )
             total = data_total
-        
+
         return total
 
     @extend_schema_field(serializers.DictField(allow_null=True))
@@ -135,17 +140,19 @@ def get_data_files_info(self, capture: Capture) -> dict[str, Any]:
     def get_center_frequency_ghz(self, capture: Capture) -> float | None:
         """Get the center frequency in GHz from the capture model property."""
         return capture.center_frequency_ghz
-    
+
     @extend_schema_field(serializers.FloatField(allow_null=True))
     def get_sample_rate_mhz(self, capture: Capture) -> float | None:
-        """Get the sample rate in MHz from the capture model property. None if not indexed in OpenSearch."""
+        """Sample rate in MHz from the model. None if not indexed in OpenSearch."""
         return capture.sample_rate_mhz
 
     @extend_schema_field(serializers.IntegerField(allow_null=True))
     def get_length_of_capture_ms(self, capture: Capture) -> int | None:
-        """Get the length of the capture in milliseconds. OpenSearch bounds are in seconds."""
+        """Capture length in milliseconds (OpenSearch bounds are seconds)."""
         try:
-            start_time, end_time = get_capture_bounds(capture.capture_type, str(capture.uuid))
+            start_time, end_time = get_capture_bounds(
+                capture.capture_type, str(capture.uuid)
+            )
             return (end_time - start_time) * 1000
         except (ValueError, IndexError, KeyError):
             return None
@@ -160,13 +167,14 @@ def get_file_cadence_ms(self, capture: Capture) -> int | None:
 
     @extend_schema_field(serializers.IntegerField(allow_null=True))
     def get_capture_start_epoch_sec(self, capture: Capture) -> int | None:
-        """Get the capture start time as Unix epoch seconds. None if not indexed in OpenSearch."""
+        """Capture start as Unix epoch seconds. None if not in OpenSearch."""
         try:
             start_time, _ = get_capture_bounds(capture.capture_type, str(capture.uuid))
-            return start_time
         except (ValueError, IndexError, KeyError):
             return None
-        
+        else:
+            return start_time
+
     @extend_schema_field(serializers.DictField)
     def get_capture_props(self, capture: Capture) -> dict[str, Any]:
         """Retrieve the indexed metadata for the capture."""
@@ -382,21 +390,25 @@ def get_total_file_size(self, obj: dict[str, Any]) -> int | None:
         """Get the total file size across all channels."""
         if obj["capture_type"] != CaptureType.DigitalRF:
             return None
-        
+
         total_size = 0
         for channel_data in obj["channels"]:
             capture_uuid = channel_data["uuid"]
             capture = Capture.objects.get(uuid=capture_uuid)
             all_files = get_capture_files(capture, include_deleted=False)
             result = all_files.aggregate(total_size=Sum("size"))
             total_size += result["total_size"] or 0
-        
+
         data_total = self.get_data_files_info(obj).get("total_size", 0)
-        
+
         if total_size < data_total:
             logging.getLogger(__name__).warning(
-                "Composite capture: total_file_size (%s) < data_files_total_size (%s); using data total.",
-                total_size, data_total,
+                (
+                    "Composite capture: total_file_size (%s) < "
+                    "data_files_total_size (%s); using data total."
+                ),
+                total_size,
+                data_total,
             )
             total_size = data_total
         return total_size
@@ -405,7 +417,7 @@ def get_data_files_info(self, obj: dict[str, Any]) -> dict[str, Any]:
         """Get the data files info for the composite capture."""
         if obj["capture_type"] != CaptureType.DigitalRF:
             return {}
-        
+
         total_count = 0
         total_size = 0
         for channel_data in obj["channels"]:
@@ -418,7 +430,9 @@ def get_data_files_info(self, obj: dict[str, Any]) -> dict[str, Any]:
         return {
             "count": total_count,
             "total_size": total_size,
-            "per_data_file_size": (float(total_size) / total_count) if total_count else None,
+            "per_data_file_size": (float(total_size) / total_count)
+            if total_count
+            else None,
         }
 
     @extend_schema_field(serializers.CharField)
@@ -464,12 +478,11 @@ def get_capture_start_epoch_sec(self, obj: dict[str, Any]) -> int | None:
             return None
         try:
             capture = Capture.objects.get(uuid=channels[0]["uuid"])
-            start_time, _ = get_capture_bounds(
-                capture.capture_type, str(capture.uuid)
-            )
-            return start_time
+            start_time, _ = get_capture_bounds(capture.capture_type, str(capture.uuid))
         except (ValueError, IndexError, KeyError):
             return None
+        else:
+            return start_time
 
 
 def build_composite_capture_data(captures: list[Capture]) -> dict[str, Any]:

gateway/sds_gateway/api_methods/tests/factories.py

@@ -13,13 +13,13 @@
 from unittest.mock import patch
 
 from django.core.files.base import ContentFile
-from faker import Faker as FakerInstance
 from factory import Faker as FactoryFaker
 from factory import LazyAttribute
 from factory import LazyFunction
-from factory import post_generation
 from factory import Sequence
+from factory import post_generation
 from factory.django import DjangoModelFactory
+from faker import Faker as FakerInstance
 
 from sds_gateway.api_methods.helpers.temporal_filtering import drf_rf_filename_from_ms
 from sds_gateway.api_methods.models import Capture
@@ -30,9 +30,10 @@
 from sds_gateway.api_methods.models import UserSharePermission
 from sds_gateway.users.tests.factories import UserFactory
 
-# Standalone Faker for LazyFunction callbacks (not factory_boy's FactoryFaker declaration)
+# Standalone Faker for LazyFunction callbacks (not FactoryFaker from factory_boy)
 _faker = FakerInstance()
 
+
 class DatasetFactory(DjangoModelFactory):
     """Factory for creating Dataset instances for testing.
 
@@ -237,26 +238,23 @@ class Meta:
 
     channel = FactoryFaker("word")
     capture_type = "drf"
-    top_level_dir = LazyFunction(
-        lambda: _faker.file_path(depth=2).replace("/", "_")
-    )
+    top_level_dir = LazyFunction(lambda: _faker.file_path(depth=2).replace("/", "_"))
     owner = FactoryFaker("subfactory", factory=UserFactory)
     name = FactoryFaker("slug")
     index_name = "captures-drf"
 
 
 class DRFDataFileFactory(DjangoModelFactory):
-    """Factory for creating DRF data file instances for testing.
-    
-    This factory creates realistic DRF data file objects that represent files stored in the system.
-    It generates test data for file metadata and creates a Django ContentFile for the actual file content.
-    
-    The factory creates files with realistic metadata including size, checksums, and proper file extensions.
-    It also handles the creation of the Django file field with test content.
+    """Factory for DRF data file instances used in tests.
+
+    Creates file metadata plus a Django ContentFile for content. Includes
+    checksums, sizes, and extensions; wires the file field for uploads.
     """
 
     uuid = FactoryFaker("uuid4")
-    directory = LazyAttribute(lambda obj: f"/files/{obj.owner.email}/{obj.capture.top_level_dir}/")
+    directory = LazyAttribute(
+        lambda obj: f"/files/{obj.owner.email}/{obj.capture.top_level_dir}/"
+    )
     name = Sequence(lambda n: drf_rf_filename_from_ms(1000 + n * 1000))
     media_type = "application/x-hdf5"
     permissions = "rw-r----"
@@ -275,8 +273,6 @@ def file(self, create, extracted, **kwargs):
         else:
             content = b"test drf file content"
             self.file = ContentFile(content, name=self.name)
-    
-
 
     class Meta:
         model = File
@@ -316,7 +312,9 @@ class UserSharePermissionFactory(DjangoModelFactory):
 
     owner = FactoryFaker("subfactory", factory=UserFactory)
     shared_with = FactoryFaker("subfactory", factory=UserFactory)
-    item_type = FactoryFaker("random_element", elements=[ItemType.DATASET, ItemType.CAPTURE])
+    item_type = FactoryFaker(
+        "random_element", elements=[ItemType.DATASET, ItemType.CAPTURE]
+    )
     item_uuid = FactoryFaker("uuid4")
     is_enabled = True
     message = FactoryFaker("sentence", nb_words=5)

gateway/sds_gateway/api_methods/tests/test_celery_tasks.py

@@ -25,6 +25,7 @@
 from sds_gateway.api_methods.models import ItemType
 from sds_gateway.api_methods.models import TemporaryZipFile
 from sds_gateway.api_methods.models import ZipFileStatus
+from sds_gateway.api_methods.tasks import _get_item_files
 from sds_gateway.api_methods.tasks import acquire_user_lock
 from sds_gateway.api_methods.tasks import check_celery_task
 from sds_gateway.api_methods.tasks import check_disk_space_available
@@ -36,7 +37,6 @@
 from sds_gateway.api_methods.tasks import get_user_task_status
 from sds_gateway.api_methods.tasks import is_user_locked
 from sds_gateway.api_methods.tasks import release_user_lock
-from sds_gateway.api_methods.tasks import _get_item_files
 from sds_gateway.api_methods.tasks import send_item_files_email
 from sds_gateway.api_methods.utils.disk_utils import estimate_disk_size
 
@@ -1270,7 +1270,7 @@ def test_get_item_files_with_temporal_bounds_returns_expected_rf_subset(self):
             "sds_gateway.api_methods.helpers.temporal_filtering.get_opensearch_client"
         ) as m:
             m.return_value.get.return_value = mock_response
-            # Relative ms: 1000–4000 from capture start → absolute 2s–5s filenames
+            # Relative ms: 1000-4000 from capture start; absolute 2s-5s filenames
             result = _get_item_files(
                 self.user,
                 self.capture,
@@ -1279,10 +1279,12 @@ def test_get_item_files_with_temporal_bounds_returns_expected_rf_subset(self):
                 end_time=4000,
             )
         names = [f.name for f in result]
-        # DRF files in [2s, 5s] inclusive (see filter_capture_data_files_selection_bounds)
+        # DRF files [2s,5s] inclusive; see filter_capture_data_files_selection_bounds
         expected_rf = [f"rf@{i}.000.h5" for i in range(2, 6)]
         rf_names = sorted(n for n in names if n.startswith("rf@"))
-        assert rf_names == expected_rf, f"Expected RF files {expected_rf}, got {rf_names}"
-        # Metadata / non-DRF capture files from setUp are still included in the download set
+        assert rf_names == expected_rf, (
+            f"Expected RF files {expected_rf}, got {rf_names}"
+        )
+        # Metadata and non-DRF files from setUp stay in the download set
         assert "test_file1.txt" in names
         assert "test_file2.txt" in names