Add isolated environment variable getters

BrianMichell · tasansal · commit 879ddca88b20 · 2025-10-25T16:00:11.000-05:00
diff --git a/src/mdio/api/_environ.py b/src/mdio/api/_environ.py
@@ -0,0 +1,121 @@
+"""Environment variable management for MDIO operations."""
+
+from os import getenv
+from typing import Final
+
+from psutil import cpu_count
+
+from mdio.converters.exceptions import EnvironmentFormatError
+
+
+class Environment:
+    """Unified API for accessing and validating MDIO environment variables."""
+
+    # Environment variable keys and defaults
+    _EXPORT_CPUS_KEY: Final[str] = "MDIO__EXPORT__CPU_COUNT"
+    _IMPORT_CPUS_KEY: Final[str] = "MDIO__IMPORT__CPU_COUNT"
+    _GRID_SPARSITY_RATIO_WARN_KEY: Final[str] = "MDIO__GRID__SPARSITY_RATIO_WARN"
+    _GRID_SPARSITY_RATIO_LIMIT_KEY: Final[str] = "MDIO__GRID__SPARSITY_RATIO_LIMIT"
+    _SAVE_SEGY_FILE_HEADER_KEY: Final[str] = "MDIO__IMPORT__SAVE_SEGY_FILE_HEADER"
+    _MDIO_SEGY_SPEC_KEY: Final[str] = "MDIO__SEGY__SPEC"
+    _RAW_HEADERS_KEY: Final[str] = "MDIO__IMPORT__RAW_HEADERS"
+    _IGNORE_CHECKS_KEY: Final[str] = "MDIO_IGNORE_CHECKS"
+    _CLOUD_NATIVE_KEY: Final[str] = "MDIO__IMPORT__CLOUD_NATIVE"
+
+    # Default values
+    _EXPORT_CPUS_DEFAULT: Final[int] = cpu_count(logical=True)
+    _IMPORT_CPUS_DEFAULT: Final[int] = cpu_count(logical=True)
+    _GRID_SPARSITY_RATIO_WARN_DEFAULT: Final[str] = "2"
+    _GRID_SPARSITY_RATIO_LIMIT_DEFAULT: Final[str] = "10"
+    _SAVE_SEGY_FILE_HEADER_DEFAULT: Final[str] = "false"
+    _MDIO_SEGY_SPEC_DEFAULT: Final[None] = None
+    _RAW_HEADERS_DEFAULT: Final[str] = "false"
+    _IGNORE_CHECKS_DEFAULT: Final[str] = "false"
+    _CLOUD_NATIVE_DEFAULT: Final[str] = "false"
+
+    @classmethod
+    def _get_env_value(cls, key: str, default: str | int | None) -> str | None:
+        """Get environment variable value with fallback to default."""
+        if isinstance(default, int):
+            default = str(default)
+        return getenv(key, default)
+
+    @staticmethod
+    def _parse_bool(value: str | None) -> bool:
+        """Parse string value to boolean."""
+        if value is None:
+            return False
+        return value.lower() in ("1", "true", "yes", "on")
+
+    @staticmethod
+    def _parse_int(value: str | None, key: str) -> int:
+        """Parse string value to integer with validation."""
+        if value is None:
+            raise EnvironmentFormatError(key, "int")
+        try:
+            return int(value)
+        except ValueError as e:
+            raise EnvironmentFormatError(key, "int") from e
+
+    @staticmethod
+    def _parse_float(value: str | None, key: str) -> float:
+        """Parse string value to float with validation."""
+        if value is None:
+            raise EnvironmentFormatError(key, "float")
+        try:
+            return float(value)
+        except ValueError as e:
+            raise EnvironmentFormatError(key, "float") from e
+
+    @classmethod
+    def export_cpus(cls) -> int:
+        """Number of CPUs to use for export operations."""
+        value = cls._get_env_value(cls._EXPORT_CPUS_KEY, cls._EXPORT_CPUS_DEFAULT)
+        return cls._parse_int(value, cls._EXPORT_CPUS_KEY)
+
+    @classmethod
+    def import_cpus(cls) -> int:
+        """Number of CPUs to use for import operations."""
+        value = cls._get_env_value(cls._IMPORT_CPUS_KEY, cls._IMPORT_CPUS_DEFAULT)
+        return cls._parse_int(value, cls._IMPORT_CPUS_KEY)
+
+    @classmethod
+    def grid_sparsity_ratio_warn(cls) -> float:
+        """Sparsity ratio threshold for warnings."""
+        value = cls._get_env_value(cls._GRID_SPARSITY_RATIO_WARN_KEY, cls._GRID_SPARSITY_RATIO_WARN_DEFAULT)
+        return cls._parse_float(value, cls._GRID_SPARSITY_RATIO_WARN_KEY)
+
+    @classmethod
+    def grid_sparsity_ratio_limit(cls) -> float:
+        """Sparsity ratio threshold for errors."""
+        value = cls._get_env_value(cls._GRID_SPARSITY_RATIO_LIMIT_KEY, cls._GRID_SPARSITY_RATIO_LIMIT_DEFAULT)
+        return cls._parse_float(value, cls._GRID_SPARSITY_RATIO_LIMIT_KEY)
+
+    @classmethod
+    def save_segy_file_header(cls) -> bool:
+        """Whether to save SEG-Y file headers."""
+        value = cls._get_env_value(cls._SAVE_SEGY_FILE_HEADER_KEY, cls._SAVE_SEGY_FILE_HEADER_DEFAULT)
+        return cls._parse_bool(value)
+
+    @classmethod
+    def mdio_segy_spec(cls) -> str | None:
+        """Path to MDIO SEG-Y specification file."""
+        return cls._get_env_value(cls._MDIO_SEGY_SPEC_KEY, cls._MDIO_SEGY_SPEC_DEFAULT)
+
+    @classmethod
+    def raw_headers(cls) -> bool:
+        """Whether to preserve raw headers."""
+        value = cls._get_env_value(cls._RAW_HEADERS_KEY, cls._RAW_HEADERS_DEFAULT)
+        return cls._parse_bool(value)
+
+    @classmethod
+    def ignore_checks(cls) -> bool:
+        """Whether to ignore validation checks."""
+        value = cls._get_env_value(cls._IGNORE_CHECKS_KEY, cls._IGNORE_CHECKS_DEFAULT)
+        return cls._parse_bool(value)
+
+    @classmethod
+    def cloud_native(cls) -> bool:
+        """Whether to use cloud-native mode for SEG-Y processing."""
+        value = cls._get_env_value(cls._CLOUD_NATIVE_KEY, cls._CLOUD_NATIVE_DEFAULT)
+        return cls._parse_bool(value)
diff --git a/src/mdio/converters/mdio.py b/src/mdio/converters/mdio.py
@@ -2,14 +2,13 @@
 
 from __future__ import annotations
 
-import os
 from tempfile import TemporaryDirectory
 from typing import TYPE_CHECKING
 
 import numpy as np
-from psutil import cpu_count
 from tqdm.dask import TqdmCallback
 
+from mdio.api._environ import Environment
 from mdio.api.io import _normalize_path
 from mdio.api.io import open_mdio
 from mdio.segy.blocked_io import to_segy
@@ -29,10 +28,6 @@
     from upath import UPath
 
 
-default_cpus = cpu_count(logical=True)
-NUM_CPUS = int(os.getenv("MDIO__EXPORT__CPU_COUNT", default_cpus))
-
-
 def mdio_to_segy(  # noqa: PLR0912, PLR0913, PLR0915
     segy_spec: SegySpec,
     input_path: UPath | Path | str,
@@ -148,7 +143,7 @@ def mdio_to_segy(  # noqa: PLR0912, PLR0913, PLR0915
             if client is not None:
                 block_records = block_records.compute()
             else:
-                block_records = block_records.compute(num_workers=NUM_CPUS)
+                block_records = block_records.compute(num_workers=Environment.export_cpus())
 
         ordered_files = [rec.path for rec in block_records.ravel() if rec != 0]
         ordered_files = [output_path] + ordered_files
diff --git a/src/mdio/converters/segy.py b/src/mdio/converters/segy.py
@@ -4,7 +4,6 @@
 
 import base64
 import logging
-import os
 from typing import TYPE_CHECKING
 
 import numpy as np
@@ -14,6 +13,7 @@
 from segy.standards.codes import MeasurementSystem as SegyMeasurementSystem
 from segy.standards.fields import binary as binary_header_fields
 
+from mdio.api._environ import Environment
 from mdio.api.io import _normalize_path
 from mdio.api.io import to_mdio
 from mdio.builder.schemas.chunk_grid import RegularChunkGrid
@@ -28,7 +28,6 @@
 from mdio.builder.schemas.v1.variable import VariableMetadata
 from mdio.builder.xarray_builder import to_xarray_dataset
 from mdio.constants import ZarrFormat
-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
@@ -92,8 +91,6 @@ def grid_density_qc(grid: Grid, num_traces: int) -> None:
     Raises:
         GridTraceSparsityError: If the sparsity ratio exceeds `MDIO__GRID__SPARSITY_RATIO_LIMIT`
             and `MDIO_IGNORE_CHECKS` is not set to a truthy value (e.g., "1", "true").
-        EnvironmentFormatError: If `MDIO__GRID__SPARSITY_RATIO_WARN` or
-            `MDIO__GRID__SPARSITY_RATIO_LIMIT` cannot be converted to a float.
     """
     # Calculate total possible traces in the grid (excluding sample dimension)
     grid_traces = np.prod(grid.shape[:-1], dtype=np.uint64)
@@ -102,20 +99,9 @@ def grid_density_qc(grid: Grid, num_traces: int) -> None:
     sparsity_ratio = float("inf") if num_traces == 0 else grid_traces / num_traces
 
     # Fetch and validate environment variables
-    warning_ratio_env = os.getenv("MDIO__GRID__SPARSITY_RATIO_WARN", "2")
-    error_ratio_env = os.getenv("MDIO__GRID__SPARSITY_RATIO_LIMIT", "10")
-    ignore_checks_env = os.getenv("MDIO_IGNORE_CHECKS", "false").lower()
-    ignore_checks = ignore_checks_env in ("1", "true", "yes", "on")
-
-    try:
-        warning_ratio = float(warning_ratio_env)
-    except ValueError as e:
-        raise EnvironmentFormatError("MDIO__GRID__SPARSITY_RATIO_WARN", "float") from e  # noqa: EM101
-
-    try:
-        error_ratio = float(error_ratio_env)
-    except ValueError as e:
-        raise EnvironmentFormatError("MDIO__GRID__SPARSITY_RATIO_LIMIT", "float") from e  # noqa: EM101
+    warning_ratio = Environment.grid_sparsity_ratio_warn()
+    error_ratio = Environment.grid_sparsity_ratio_limit()
+    ignore_checks = Environment.ignore_checks()
 
     # Check sparsity
     should_warn = sparsity_ratio > warning_ratio
@@ -373,7 +359,7 @@ def _populate_coordinates(
 
 
 def _add_segy_file_headers(xr_dataset: xr_Dataset, segy_file_info: SegyFileInfo) -> xr_Dataset:
-    save_file_header = os.getenv("MDIO__IMPORT__SAVE_SEGY_FILE_HEADER", "") in ("1", "true", "yes", "on")
+    save_file_header = Environment.save_segy_file_header()
     if not save_file_header:
         return xr_dataset
 
@@ -398,7 +384,7 @@ def _add_segy_file_headers(xr_dataset: xr_Dataset, segy_file_info: SegyFileInfo)
             "binaryHeader": segy_file_info.binary_header_dict,
         }
     )
-    if os.getenv("MDIO__IMPORT__RAW_HEADERS") in ("1", "true", "yes", "on"):
+    if Environment.raw_headers():
         raw_binary_base64 = base64.b64encode(segy_file_info.raw_binary_headers).decode("ascii")
         xr_dataset["segy_file_header"].attrs.update({"rawBinaryHeader": raw_binary_base64})
 
@@ -565,7 +551,7 @@ def segy_to_mdio(  # noqa PLR0913
     _, non_dim_coords = _get_coordinates(grid, segy_headers, mdio_template)
     header_dtype = to_structured_type(segy_spec.trace.header.dtype)
 
-    if os.getenv("MDIO__IMPORT__RAW_HEADERS") in ("1", "true", "yes", "on"):
+    if Environment.raw_headers():
         if zarr.config.get("default_zarr_format") == ZarrFormat.V2:
             logger.warning("Raw headers are only supported for Zarr v3. Skipping raw headers.")
         else:
diff --git a/src/mdio/segy/_workers.py b/src/mdio/segy/_workers.py
@@ -3,12 +3,12 @@
 from __future__ import annotations
 
 import logging
-import os
 from typing import TYPE_CHECKING
 
 import numpy as np
 from segy.arrays import HeaderArray
 
+from mdio.api._environ import Environment
 from mdio.api.io import _normalize_storage_options
 from mdio.segy._raw_trace_wrapper import SegyFileRawTraceWrapper
 from mdio.segy.file import SegyFileArguments
@@ -50,12 +50,7 @@ def header_scan_worker(
 
     slice_ = slice(*trace_range)
 
-    cloud_native_mode = os.getenv("MDIO__IMPORT__CLOUD_NATIVE", default="False")
-
-    if cloud_native_mode.lower() in {"true", "1"}:
-        trace_header = segy_file.trace[slice_].header
-    else:
-        trace_header = segy_file.header[slice_]
+    trace_header = segy_file.trace[slice_].header if Environment.cloud_native() else segy_file.header[slice_]
 
     if subset is not None:
         # struct field selection needs a list, not a tuple; a subset is a tuple from the template.
diff --git a/src/mdio/segy/blocked_io.py b/src/mdio/segy/blocked_io.py
@@ -3,7 +3,6 @@
 from __future__ import annotations
 
 import multiprocessing as mp
-import os
 from concurrent.futures import ProcessPoolExecutor
 from concurrent.futures import as_completed
 from pathlib import Path
@@ -13,10 +12,10 @@
 import zarr
 from dask.array import Array
 from dask.array import map_blocks
-from psutil import cpu_count
 from tqdm.auto import tqdm
 from zarr import open_group as zarr_open_group
 
+from mdio.api._environ import Environment
 from mdio.api.io import _normalize_storage_options
 from mdio.builder.schemas.v1.stats import CenteredBinHistogram
 from mdio.builder.schemas.v1.stats import SummaryStatistics
@@ -37,8 +36,6 @@
 
     from mdio.segy.file import SegyFileArguments
 
-default_cpus = cpu_count(logical=True)
-
 
 def _create_stats() -> SummaryStatistics:
     histogram = CenteredBinHistogram(bin_centers=[], counts=[])
@@ -83,8 +80,7 @@ def to_zarr(  # noqa: PLR0913, PLR0915
 
     # For Unix async writes with s3fs/fsspec & multiprocessing, use 'spawn' instead of default
     # 'fork' to avoid deadlocks on cloud stores. Slower but necessary. Default on Windows.
-    num_cpus = int(os.getenv("MDIO__IMPORT__CPU_COUNT", default_cpus))
-    num_workers = min(num_chunks, num_cpus)
+    num_workers = min(num_chunks, Environment.import_cpus())
     context = mp.get_context("spawn")
     executor = ProcessPoolExecutor(max_workers=num_workers, mp_context=context)
 
diff --git a/src/mdio/segy/compat.py b/src/mdio/segy/compat.py
@@ -8,7 +8,6 @@
 from __future__ import annotations
 
 import logging
-import os
 from importlib import metadata
 
 from packaging import version
@@ -23,6 +22,7 @@
 from segy.schema import TraceSpec
 from segy.standards.fields import binary
 
+from mdio.api._environ import Environment
 from mdio.exceptions import InvalidMDIOError
 
 MDIO_VERSION = metadata.version("multidimio")
@@ -74,7 +74,7 @@ def get_trace_fields(version_str: str) -> list[HeaderField]:
 
 def mdio_segy_spec(version_str: str | None = None) -> SegySpec:
     """Get a SEG-Y encoding spec for MDIO based on version."""
-    spec_override = os.getenv("MDIO__SEGY__SPEC")
+    spec_override = Environment.mdio_segy_spec()
 
     if spec_override is not None:
         return SegySpec.model_validate_json(spec_override)
diff --git a/src/mdio/segy/parsers.py b/src/mdio/segy/parsers.py
@@ -3,25 +3,22 @@
 from __future__ import annotations
 
 import multiprocessing as mp
-import os
 from concurrent.futures import ProcessPoolExecutor
 from itertools import repeat
 from math import ceil
 from typing import TYPE_CHECKING
 
 import numpy as np
-from psutil import cpu_count
 from tqdm.auto import tqdm
 
+from mdio.api._environ import Environment
 from mdio.segy._workers import header_scan_worker
 
 if TYPE_CHECKING:
     from segy.arrays import HeaderArray
 
     from mdio.segy.file import SegyFileArguments
 
-default_cpus = cpu_count(logical=True)
-
 
 def parse_headers(
     segy_file_kwargs: SegyFileArguments,
@@ -53,8 +50,7 @@ def parse_headers(
 
         trace_ranges.append((start, stop))
 
-    num_cpus = int(os.getenv("MDIO__IMPORT__CPU_COUNT", default_cpus))
-    num_workers = min(n_blocks, num_cpus)
+    num_workers = min(n_blocks, Environment.import_cpus())
 
     tqdm_kw = {"unit": "block", "dynamic_ncols": True}
     # For Unix async writes with s3fs/fsspec & multiprocessing, use 'spawn' instead of default
diff --git a/tests/unit/test_environment.py b/tests/unit/test_environment.py
