Pre-commit

Author: BrianMichell

pyproject.toml

@@ -208,4 +208,4 @@ module-name = "mdio"
 
 [build-system]
 requires = ["uv_build>=0.8.17,<0.9.0"]
-build-backend = "uv_build"
+build-backend = "uv_build"

src/mdio/converters/segy.py

@@ -612,4 +612,4 @@ def segy_to_mdio(  # noqa PLR0913
                 "checksum_scope": "full_file",
                 "checksum_library": "google-crc32c",
             }
-        )
+        )

src/mdio/segy/checksum.py

@@ -15,25 +15,31 @@
 from __future__ import annotations
 
 import logging
+import multiprocessing as mp
 import os
+from concurrent.futures import ProcessPoolExecutor
+from itertools import repeat
+from math import ceil
 from typing import TYPE_CHECKING
 from typing import Any
+
 import numpy as np
+from psutil import cpu_count
 from segy import SegyFile
-from mdio.segy._raw_trace_wrapper import SegyFileRawTraceWrapper
-from math import ceil
-from concurrent.futures import ProcessPoolExecutor
-from itertools import repeat
+from segy.arrays import HeaderArray
 from tqdm.auto import tqdm
 from upath import UPath
-from psutil import cpu_count
-import multiprocessing as mp
 
+from mdio.segy._raw_trace_wrapper import SegyFileRawTraceWrapper
 
-from segy.arrays import HeaderArray
 if TYPE_CHECKING:
     from mdio.segy._workers import SegyFileArguments
 
+    try:
+        from crc32c_dist_rs import DistributedCRC32C
+    except ImportError:
+        DistributedCRC32C = Any  # type: ignore[assignment,misc]
+
 default_cpus = cpu_count(logical=True)
 
 logger = logging.getLogger(__name__)
@@ -58,7 +64,7 @@
 
 
 def header_scan_worker(
-    segy_file_kwargs: "SegyFileArguments",
+    segy_file_kwargs: SegyFileArguments,
     trace_range: tuple[int, int],
     subset: list[str] | None = None,
 ) -> HeaderArray | tuple[HeaderArray, tuple[int, int, int]]:
@@ -103,7 +109,6 @@ def header_scan_worker(
     # (singleton) so we can concat and assign stuff later.
     trace_header = np.array(trace_header, dtype=new_dtype, ndmin=1)
 
-
     # Calculate checksum from the raw bytes ALREADY IN MEMORY (NO ADDITIONAL I/O!)
     raw_bytes = traces.trace_buffer_array.tobytes()
     partial_crc32c = calculate_bytes_crc32c(raw_bytes)
@@ -122,13 +127,13 @@ def header_scan_worker(
 
     return HeaderArray(trace_header), checksum_info
 
+
 def parse_headers(  # noqa: PLR0913
     segy_file_kwargs: SegyFileArguments,
     num_traces: int,
     subset: list[str] | None = None,
     block_size: int = 10000,
     progress_bar: bool = True,
-    calculate_checksum: bool = True,
 ) -> HeaderArray | tuple[HeaderArray, int]:
     """Read and parse given `byte_locations` from SEG-Y file.
 
@@ -200,9 +205,7 @@ def parse_headers(  # noqa: PLR0913
         )
 
         if progress_bar is True:
-            desc = (
-                "Scanning SEG-Y & calculating checksum"
-            )
+            desc = "Scanning SEG-Y & calculating checksum"
             lazy_work = tqdm(
                 iterable=lazy_work,
                 total=n_blocks,
@@ -215,6 +218,7 @@ def parse_headers(  # noqa: PLR0913
 
     return finalize_distributed_checksum(results, combiner)
 
+
 def is_checksum_available() -> bool:
     """Check if checksum libraries are available.
 
@@ -273,7 +277,7 @@ def calculate_bytes_crc32c(data: bytes) -> int:
     return int.from_bytes(crc.digest(), byteorder="big")
 
 
-def create_distributed_crc32c(initial_bytes: bytes, total_length: int) -> Any:
+def create_distributed_crc32c(initial_bytes: bytes, total_length: int) -> DistributedCRC32C:
     """Create a distributed CRC32C combiner instance.
 
     Args:
@@ -292,7 +296,7 @@ def create_distributed_crc32c(initial_bytes: bytes, total_length: int) -> Any:
 
 
 def finalize_distributed_checksum(
-    results: list[tuple[HeaderArray, tuple[int, int, int]]], combiner: Any
+    results: list[tuple[HeaderArray, tuple[int, int, int]]], combiner: DistributedCRC32C
 ) -> tuple[HeaderArray, int]:
     """Finalize a distributed CRC32C checksum from scan results.
 
@@ -309,8 +313,6 @@ def finalize_distributed_checksum(
     """
     require_checksum_libraries()
 
-    import numpy as np
-
     headers: list[HeaderArray] = []
     for result in results:
         headers.append(result[0])
@@ -334,4 +336,3 @@ def finalize_distributed_checksum(
     "create_distributed_crc32c",
     "finalize_distributed_checksum",
 ]
-

src/mdio/segy/parsers.py

@@ -22,6 +22,7 @@
 
 default_cpus = cpu_count(logical=True)
 
+
 def parse_headers(
     segy_file_kwargs: SegyFileArguments,
     num_traces: int,
@@ -74,4 +75,4 @@ def parse_headers(
         headers: list[HeaderArray] = list(lazy_work)
 
     # Merge blocks before return
-    return np.concatenate(headers)
+    return np.concatenate(headers)

src/mdio/segy/utilities.py

@@ -11,10 +11,9 @@
 from dask.array.core import normalize_chunks
 
 from mdio.core import Dimension
-from mdio.segy.geometry import GridOverrider
-
-from mdio.segy.checksum import should_calculate_checksum
 from mdio.segy.checksum import is_checksum_available
+from mdio.segy.checksum import should_calculate_checksum
+from mdio.segy.geometry import GridOverrider
 
 if TYPE_CHECKING:
     from numpy.typing import DTypeLike
@@ -61,10 +60,11 @@ def get_grid_plan(  # noqa:  C901, PLR0913
     """
     calculate_checksum = False
     if should_calculate_checksum() and is_checksum_available():
-        from mdio.segy.checksum import parse_headers
+        from mdio.segy.checksum import parse_headers  # noqa: PLC0415
+
         calculate_checksum = True
     else:
-        from mdio.segy.parsers import parse_headers
+        from mdio.segy.parsers import parse_headers  # noqa: PLC0415
 
     if grid_overrides is None:
         grid_overrides = {}

tests/integration/test_crc32c_checksum.py

@@ -47,7 +47,7 @@ def get_expected_crc32c(segy_path: Path) -> int:
         CRC32C checksum as integer
     """
     # Import here to keep it test-only
-    import google_crc32c
+    import google_crc32c  # noqa: PLC0415
 
     crc = google_crc32c.Checksum()
 
@@ -175,4 +175,4 @@ def test_corrupted_data_detects_changes(self, tmp_path: Path, monkeypatch: pytes
         assert mdio_crc == corrupted_crc, (
             f"MDIO checksum doesn't match corrupted file checksum. "
             f"MDIO: 0x{mdio_crc:08x}, Expected: 0x{corrupted_crc:08x}"
-        )
+        )