Manual reconciliation of merging v1_backup into v1 branch

Author: BrianMichell

src/mdio/__main__.py

@@ -7,6 +7,9 @@
 from importlib import metadata
 from pathlib import Path
 
+from typing import Any
+from typing import Callable
+
 import click
 
 
@@ -30,13 +33,13 @@ class MyCLI(click.MultiCommand):
     - plugin_folder: Path to the directory containing command modules.
     """
 
-    def __init__(self, plugin_folder: Path, *args, **kwargs):
+    def __init__(self, plugin_folder: Path, *args: Any, **kwargs: Any) -> None:  # noqa: ANN401
         """Initializer function."""
         super().__init__(*args, **kwargs)
         self.plugin_folder = plugin_folder
         self.known_modules = KNOWN_MODULES
 
-    def list_commands(self, ctx: click.Context) -> list[str]:
+    def list_commands(self, _ctx: click.Context) -> list[str]:
         """List commands available under `commands` module."""
         rv = []
         for filename in self.plugin_folder.iterdir():
@@ -47,7 +50,7 @@ def list_commands(self, ctx: click.Context) -> list[str]:
         rv.sort()
         return rv
 
-    def get_command(self, ctx: click.Context, name: str) -> Callable | None:
+    def get_command(self, _ctx: click.Context, name: str) -> Callable | None:
         """Get command implementation from `commands` module."""
         try:
             filepath = self.plugin_folder / f"{name}.py"

src/mdio/api/accessor.py

@@ -2,9 +2,10 @@
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 import logging
 
-import dask.array as da
 import numpy as np
 import numpy.typing as npt
 import zarr
@@ -18,6 +19,9 @@
 from mdio.core.exceptions import MDIONotFoundError
 from mdio.exceptions import ShapeError
 
+if TYPE_CHECKING:
+    import dask.array as da
+    from numpy.typing import NDArray
 
 logger = logging.getLogger(__name__)
 
@@ -181,7 +185,7 @@ def __init__(
         self._set_attributes()
         self._open_arrays()
 
-    def _validate_store(self, storage_options):
+    def _validate_store(self, storage_options: dict[str, str] | None) -> None:
         """Method to validate the provided store."""
         if storage_options is None:
             storage_options = {}
@@ -194,7 +198,7 @@ def _validate_store(self, storage_options):
             disk_cache=self._disk_cache,
         )
 
-    def _connect(self):
+    def _connect(self) -> None:
         """Open the zarr root."""
         try:
             if self.mode in {"r", "r+"}:
@@ -212,11 +216,11 @@ def _connect(self):
             )
             raise MDIONotFoundError(msg) from e
 
-    def _deserialize_grid(self):
+    def _deserialize_grid(self) -> None:
         """Deserialize grid from Zarr metadata."""
         self.grid = Grid.from_zarr(self.root)
 
-    def _set_attributes(self):
+    def _set_attributes(self) -> None:
         """Deserialize attributes from Zarr metadata."""
         self.trace_count = self.root.attrs["trace_count"]
         self.stats = {
@@ -231,7 +235,7 @@ def _set_attributes(self):
         self.n_dim = len(self.shape)
 
         # Access pattern attributes
-        data_array_name = "_".join(["chunked", self.access_pattern])
+        data_array_name = f"chunked_{self.access_pattern}"
         self.chunks = self._data_group[data_array_name].chunks
         self._orig_chunks = self.chunks
 
@@ -251,15 +255,15 @@ def _set_attributes(self):
             self._orig_chunks = self.chunks
             self.chunks = new_chunks
 
-    def _open_arrays(self):
+    def _open_arrays(self) -> None:
         """Open arrays with requested backend."""
-        data_array_name = "_".join(["chunked", self.access_pattern])
-        header_array_name = "_".join(["chunked", self.access_pattern, "trace_headers"])
+        data_array_name = f"chunked_{self.access_pattern}"
+        header_array_name = f"chunked_{self.access_pattern}_trace_headers"
 
-        trace_kwargs = dict(
-            group_handle=self._data_group,
-            name=data_array_name,
-        )
+        trace_kwargs = {
+            "group_handle": self._data_group,
+            "name": data_array_name,
+        }
 
         if self._backend == "dask":
             trace_kwargs["chunks"] = self.chunks
@@ -271,10 +275,10 @@ def _open_arrays(self):
             logger.info(f"Setting MDIO in-memory chunks to {dask_chunks}")
             self.chunks = dask_chunks
 
-        header_kwargs = dict(
-            group_handle=self._metadata_group,
-            name=header_array_name,
-        )
+        header_kwargs = {
+            "group_handle": self._metadata_group,
+            "name": header_array_name,
+        }
 
         if self._backend == "dask":
             header_kwargs["chunks"] = self.chunks[:-1]
@@ -406,7 +410,7 @@ def __setitem__(self, key: int | tuple, value: npt.ArrayLike) -> None:
 
     def coord_to_index(
         self,
-        *args,
+        *args: int | list[int],
         dimensions: str | list[str] | None = None,
     ) -> tuple[NDArray[int], ...]:
         """Convert dimension coordinate to zero-based index.
@@ -437,6 +441,7 @@ def coord_to_index(
             to indicies of that dimension
 
         Raises:
+            KeyError: if a requested dimension doesn't exist.
             ShapeError: if number of queries don't match requested dimensions.
             ValueError: if requested coordinates don't exist.
 
@@ -490,10 +495,15 @@ def coord_to_index(
         if dimensions is None:
             dims = self.grid.dims
         else:
-            dims = [self.grid.select_dim(dim_name) for dim_name in dimensions]
-
-        dim_indices = tuple()
-        for mdio_dim, dim_query_coords in zip(dims, queries):  # noqa: B905
+            for query_dim in dimensions:
+                try:
+                    dims.append(self.grid.select_dim(query_dim))
+                except ValueError as err:
+                    msg = f"Requested dimension {query_dim} does not exist."
+                    raise KeyError(msg) from err
+
+        dim_indices = ()
+        for mdio_dim, dim_query_coords in zip(dims, queries):
             # Make sure all coordinates exist.
             query_diff = np.setdiff1d(dim_query_coords, mdio_dim.coords)
             if len(query_diff) > 0:
@@ -510,14 +520,14 @@ def coord_to_index(
 
         return dim_indices if len(dim_indices) > 1 else dim_indices[0]
 
-    def copy(
+    def copy(  # noqa: PLR0913
         self,
         dest_path_or_buffer: str,
         excludes: str = "",
         includes: str = "",
         storage_options: dict | None = None,
         overwrite: bool = False,
-    ):
+    ) -> None:
         """Makes a copy of an MDIO file with or without all arrays.
 
         Refer to mdio.api.convenience.copy for full documentation.
@@ -576,7 +586,7 @@ class MDIOReader(MDIOAccessor):
             `fsspec` documentation for more details.
     """
 
-    def __init__(
+    def __init__(  # noqa: PLR0913
         self,
         mdio_path_or_buffer: str,
         access_pattern: str = "012",
@@ -632,7 +642,7 @@ class MDIOWriter(MDIOAccessor):
             `fsspec` documentation for more details.
     """
 
-    def __init__(
+    def __init__(  # noqa: PLR0913
         self,
         mdio_path_or_buffer: str,
         access_pattern: str = "012",

src/mdio/api/convenience.py

@@ -13,6 +13,9 @@
 
 
 if TYPE_CHECKING:
+    from pathlib import Path
+    from typing import Any
+
     from numcodecs.abc import Codec
     from numpy.typing import NDArray
     from zarr import Array
@@ -23,10 +26,10 @@
 
 def copy_mdio(  # noqa: PLR0913
     source: MDIOReader,
-    dest_path_or_buffer: str,
+    dest_path_or_buffer: str | Path,
     excludes: str = "",
     includes: str = "",
-    storage_options: dict | None = None,
+    storage_options: dict[str, Any] | None = None,
     overwrite: bool = False,
 ) -> None:
     """Copy MDIO file.
@@ -135,7 +138,8 @@ def create_rechunk_plan(
 
     for chunks, suffix in zip(chunks_list, suffix_list):  # noqa: B905
         norm_chunks = [
-            min(chunk, size) for chunk, size in zip(chunks, source.shape)  # noqa: B905
+            min(chunk, size)
+              for chunk, size in zip(chunks, source.shape)  # noqa: B905
         ]
 
         if suffix == source.access_pattern:

src/mdio/api/io_utils.py

@@ -2,16 +2,20 @@
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
 from typing import Any
 
 import dask.array as da
 import zarr
 from zarr.storage import FSStore
 
+if TYPE_CHECKING:
+    from pathlib import Path
+
 
 def process_url(
-    url: str,
-    mode: str,
+    url: str | Path,
+    mode: str | Path,
     storage_options: dict[str, Any],
     memory_cache_size: int,
     disk_cache: bool,
@@ -82,20 +86,20 @@ def process_url(
         ... )
     """
     if disk_cache is True:
-        url = "::".join(["simplecache", url])
+        url = f"simplecache::{url}"
 
     # Strip whitespaces and slashes from end of string
-    url = url.rstrip("/ ")
+    url_str = str(url).rstrip("/ ")
 
     # Flag for checking write access
-    check = True if mode == "w" else False
+    # check = True if mode == "w" else False
 
     # TODO: Turning off write checking now because zarr has a bug.
     #  Get rid of this once bug is fixed.
     check = False
 
     store = FSStore(
-        url=url,
+        url=url_str,
         check=check,
         create=check,
         mode=mode,
@@ -123,7 +127,9 @@ def open_zarr_array(group_handle: zarr.Group, name: str) -> zarr.Array:
     return group_handle[name]
 
 
-def open_zarr_array_dask(group_handle: zarr.Group, name: str, **kwargs) -> da.Array:
+def open_zarr_array_dask(
+    group_handle: zarr.Group, name: str, **kwargs: dict[str, Any]
+) -> da.Array:
     """Open Zarr array lazily using Dask.
 
     Note: All other kwargs get passed to dask.array.from_zarr()

src/mdio/commands/copy.py

@@ -55,7 +55,7 @@
     help="Flag to overwrite if mdio file if it exists",
     show_default=True,
 )
-def copy(
+def copy(  # noqa: PLR0913
     source_mdio_path: str,
     target_mdio_path: str,
     access_pattern: str = "012",

src/mdio/commands/segy.py

@@ -393,7 +393,7 @@ def segy_import(  # noqa: PLR0913
     show_default=True,
     show_choices=True,
 )
-def segy_export(
+def segy_export(  # noqa: PLR0913
     mdio_file: str,
     segy_path: str,
     access_pattern: str,

src/mdio/converters/exceptions.py

@@ -4,10 +4,10 @@
 class EnvironmentFormatError(Exception):
     """Raised when environment variable is of the wrong format."""
 
-    def __init__(self, name, format, msg: str = ""):
+    def __init__(self, name: str, format_: str, msg: str = "") -> None:
         """Initialize error."""
         self.message = (
-            f"Environment variable: {name} not of expected format: {format}. "
+            f"Environment variable: {name} not of expected format: {format_}. "
         )
         self.message += f"\n{msg}" if msg else ""
         super().__init__(self.message)
@@ -16,7 +16,7 @@ def __init__(self, name, format, msg: str = ""):
 class GridTraceCountError(Exception):
     """Raised when grid trace counts don't match the SEG-Y trace count."""
 
-    def __init__(self, grid_traces, segy_traces):
+    def __init__(self, grid_traces: int, segy_traces: int) -> None:
         """Initialize error."""
         self.message = (
             f"{grid_traces} != {segy_traces}"
@@ -32,7 +32,7 @@ def __init__(self, grid_traces, segy_traces):
 class GridTraceSparsityError(Exception):
     """Raised when mdio grid will be sparsely populated from SEG-Y traces."""
 
-    def __init__(self, shape, num_traces, msg: str = ""):
+    def __init__(self, shape: tuple[int, ...], num_traces: int, msg: str = "") -> None:
         """Initialize error."""
         self.message = (
             f"Grid shape: {shape} but SEG-Y tracecount: {num_traces}. "

src/mdio/converters/mdio.py

@@ -4,6 +4,7 @@
 
 import os
 from os import path
+from pathlib import Path
 from tempfile import TemporaryDirectory
 
 import numpy as np
@@ -28,8 +29,8 @@
 
 
 def mdio_to_segy(  # noqa: C901
-    mdio_path_or_buffer: str,
-    output_segy_path: str,
+    mdio_path_or_buffer: str | Path,
+    output_segy_path: str | Path,
     endian: str = "big",
     access_pattern: str = "012",
     storage_options: dict = None,