Linting

Author: BrianMichell

src/mdio/converters/segy.py

@@ -14,19 +14,20 @@
 
 from mdio.api.io import _normalize_path
 from mdio.api.io import to_mdio
+from mdio.builder.schemas.chunk_grid import RegularChunkGrid
+from mdio.builder.schemas.chunk_grid import RegularChunkShape
+from mdio.builder.schemas.compressors import Blosc
+from mdio.builder.schemas.compressors import BloscCname
+from mdio.builder.schemas.dtype import ScalarType
 from mdio.builder.schemas.v1.units import LengthUnitEnum
 from mdio.builder.schemas.v1.units import LengthUnitModel
+from mdio.builder.schemas.v1.variable import VariableMetadata
 from mdio.builder.xarray_builder import to_xarray_dataset
 from mdio.converters.exceptions import EnvironmentFormatError
 from mdio.converters.exceptions import GridTraceCountError
 from mdio.converters.exceptions import GridTraceSparsityError
 from mdio.converters.type_converter import to_structured_type
 from mdio.core.grid import Grid
-from mdio.builder.schemas.chunk_grid import RegularChunkGrid
-from mdio.builder.schemas.chunk_grid import RegularChunkShape
-from mdio.builder.schemas.compressors import Blosc
-from mdio.builder.schemas.compressors import BloscCname
-from mdio.builder.schemas.dtype import ScalarType
 from mdio.segy import blocked_io
 from mdio.segy.utilities import get_grid_plan
 
@@ -340,6 +341,7 @@ def _add_grid_override_to_metadata(dataset: Dataset, grid_overrides: dict[str, A
 
 def _add_raw_headers_to_template(mdio_template: AbstractDatasetTemplate) -> AbstractDatasetTemplate:
     """Add raw headers capability to the MDIO template by monkey-patching its _add_variables method.
+
     This function modifies the template's _add_variables method to also add a raw headers variable
     with the following characteristics:
     - Same rank as the Headers variable (all dimensions except vertical)
@@ -351,9 +353,11 @@ def _add_raw_headers_to_template(mdio_template: AbstractDatasetTemplate) -> Abst
     - Chunked the same as the Headers variable
     Args:
         mdio_template: The MDIO template to mutate
+    Returns:
+        The mutated MDIO template
     """
     # Check if raw headers enhancement has already been applied to avoid duplicate additions
-    if hasattr(mdio_template, '_mdio_raw_headers_enhanced'):
+    if hasattr(mdio_template, "_mdio_raw_headers_enhanced"):
         return mdio_template
 
     # Store the original _add_variables method
@@ -368,7 +372,6 @@ def enhanced_add_variables() -> None:
 
         # Create chunk grid metadata
         chunk_metadata = RegularChunkGrid(configuration=RegularChunkShape(chunk_shape=chunk_shape))
-        from mdio.builder.schemas.v1.variable import VariableMetadata
 
         # Add the raw headers variable using the builder's add_variable method
         mdio_template._builder.add_variable(

src/mdio/segy/_workers.py

@@ -154,39 +154,39 @@ def trace_worker(  # noqa: PLR0913
         )
     if raw_header_key in worker_variables:
         tmp_raw_headers = np.zeros_like(dataset[raw_header_key])
-        
+
         # Get the indices where we need to place results
         live_mask = not_null
         live_positions = np.where(live_mask.ravel())[0]
-        
+
         if len(live_positions) > 0:
             # Calculate byte ranges for headers
-            HEADER_SIZE = 240
+            header_size = 240
             trace_offset = segy_file.spec.trace.offset
             trace_itemsize = segy_file.spec.trace.itemsize
-            
+
             starts = []
             ends = []
             for global_trace_idx in live_trace_indexes:
                 header_start = trace_offset + global_trace_idx * trace_itemsize
-                header_end = header_start + HEADER_SIZE
+                header_end = header_start + header_size
                 starts.append(header_start)
                 ends.append(header_end)
-            
+
             # Capture raw bytes
             raw_header_bytes = merge_cat_file(segy_file.fs, segy_file.url, starts, ends)
-            
+
             # Convert and place results
             raw_headers_array = np.frombuffer(bytes(raw_header_bytes), dtype="|V240")
             tmp_raw_headers.ravel()[live_positions] = raw_headers_array
-        
+
         ds_to_write[raw_header_key] = Variable(
             ds_to_write[raw_header_key].dims,
             tmp_raw_headers,
             attrs=ds_to_write[raw_header_key].attrs,
             encoding=ds_to_write[raw_header_key].encoding,
         )
-    
+
     data_variable = ds_to_write[data_variable_name]
     fill_value = _get_fill_value(ScalarType(data_variable.dtype.name))
     tmp_samples = np.full_like(data_variable, fill_value=fill_value)