Cleanup

Author: BrianMichell

src/mdio/segy/_disaster_recovery_wrapper.py

@@ -2,111 +2,66 @@
 
 from __future__ import annotations
 
-import numpy as np
 from typing import TYPE_CHECKING
-from segy.transforms import ByteSwapTransform
-from segy.transforms import IbmFloatTransform
 
 if TYPE_CHECKING:
+    import numpy as np
     from segy.file import SegyFile
-    from segy.indexing import HeaderIndexer
-    from segy.transforms import Transform, TransformPipeline, ByteSwapTransform, IbmFloatTransform
+    from segy.transforms import Transform, ByteSwapTransform, IbmFloatTransform
     from numpy.typing import NDArray
 
 def _reverse_single_transform(data: NDArray, transform: Transform) -> NDArray:
-    """Reverse a single transform operation.
-
-    Args:
-        data: The data to reverse transform
-        transform: The transform to reverse
-
-    Returns:
-        Data with the transform reversed
-    """
-    # Import here to avoid circular imports
-    from segy.transforms import get_endianness
+    """Reverse a single transform operation."""
     from segy.schema import Endianness
 
     if isinstance(transform, ByteSwapTransform):
-        # For byte swap, we need to reverse the endianness conversion
-        # If the transform was converting to little-endian, we need to convert back to big-endian
-
-        # If transform was converting TO little-endian, we need to convert TO big-endian
+        # Reverse the endianness conversion
         # TODO: I don't think this is correct
         if transform.target_order == Endianness.LITTLE:
             reverse_target = Endianness.BIG
         else:
             reverse_target = Endianness.LITTLE
 
         reverse_transform = ByteSwapTransform(reverse_target)
-        result = reverse_transform.apply(data)
-
-        return result
+        return reverse_transform.apply(data)
 
     elif isinstance(transform, IbmFloatTransform):
-        # Reverse IBM float conversion by swapping direction
+        # Reverse IBM float conversion
         reverse_direction = "to_ibm" if transform.direction == "to_ieee" else "to_ieee"
         reverse_transform = IbmFloatTransform(reverse_direction, transform.keys)
         return reverse_transform.apply(data)
 
     else:
         # For unknown transforms, return data unchanged
-        # This maintains compatibility if new transforms are added
         return data
 
 def get_header_raw_and_transformed(
     segy_file: SegyFile,
-    indices: int | list[int] | np.ndarray | slice
+    indices: int | list[int] | NDArray | slice
 ) -> tuple[NDArray, NDArray, NDArray]:
-    """Convenience function to get both raw and transformed header data.
-
-    This is a drop-in replacement that provides the functionality you requested
-    without modifying the segy package.
+    """Get both raw and transformed header data.
 
     Args:
         segy_file: The SegyFile instance
         indices: Which headers to retrieve
 
     Returns:
-        Tuple of (raw_headers, transformed_headers)
-
-    Example:
-        from header_raw_transformed_accessor import get_header_raw_and_transformed
-
-        # Single header
-        raw_hdr, transformed_hdr = get_header_raw_and_transformed(segy_file, 0)
-
-        # Multiple headers
-        raw_hdrs, transformed_hdrs = get_header_raw_and_transformed(segy_file, [0, 1, 2])
-
-        # Slice of headers
-        raw_hdrs, transformed_hdrs = get_header_raw_and_transformed(segy_file, slice(0, 10))
+        Tuple of (raw_headers, transformed_headers, traces)
     """
 
     traces = segy_file.trace[indices]
-
     transformed_headers = traces.header
 
-    # Reverse the transforms on the already-loaded transformed data
-    # This eliminates the second disk read entirely!
+    # Reverse transforms to get raw data
     raw_headers = _reverse_transforms(transformed_headers, segy_file.header.transform_pipeline)
 
     return raw_headers, transformed_headers, traces
 
 def _reverse_transforms(transformed_data: NDArray, transform_pipeline) -> NDArray:
-    """Reverse the transform pipeline to get raw data from transformed data.
-
-    Args:
-        transformed_data: Data that has been processed through the transform pipeline
-        transform_pipeline: The transform pipeline to reverse
-
-    Returns:
-        Raw data equivalent to what was read directly from filesystem
-    """
-    # Start with the transformed data
+    """Reverse the transform pipeline to get raw data."""
     raw_data = transformed_data.copy() if hasattr(transformed_data, 'copy') else transformed_data
 
-    # Apply transforms in reverse order with reversed operations
+    # Apply transforms in reverse order
     for transform in reversed(transform_pipeline.transforms):
         raw_data = _reverse_single_transform(raw_data, transform)
 

src/mdio/segy/_workers.py

@@ -121,7 +121,6 @@ def trace_worker(  # noqa: PLR0913
     zarr_config.set({"threading.max_workers": 1})
 
     live_trace_indexes = local_grid_map[not_null].tolist()
-    # traces = segy_file.trace[live_trace_indexes]
     raw_headers, transformed_headers, traces = get_header_raw_and_transformed(segy_file, live_trace_indexes)
 
     header_key = "headers"
@@ -136,14 +135,11 @@ def trace_worker(  # noqa: PLR0913
         worker_variables.append(raw_header_key)
 
     ds_to_write = dataset[worker_variables]
-    # raw_headers, transformed_headers = get_header_raw_and_transformed(segy_file, live_trace_indexes)
 
     if header_key in worker_variables:
         # Create temporary array for headers with the correct shape
-        # TODO(BrianMichell): Implement this better so that we can enable fill values without changing the code. #noqa: TD003
         tmp_headers = np.zeros_like(dataset[header_key])
         tmp_headers[not_null] = transformed_headers
-        # tmp_headers[not_null] = traces.header
         # Create a new Variable object to avoid copying the temporary array
         # The ideal solution is to use `ds_to_write[header_key][:] = tmp_headers`
         # but Xarray appears to be copying memory instead of doing direct assignment.