Refactor environment managmenet to expose basic functions for lighter imports

Author: BrianMichell

src/mdio/api/_environ.py

@@ -1,121 +1,116 @@
 """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
 
+# Environment variable keys
+_EXPORT_CPUS_KEY = "MDIO__EXPORT__CPU_COUNT"
+_IMPORT_CPUS_KEY = "MDIO__IMPORT__CPU_COUNT"
+_GRID_SPARSITY_RATIO_WARN_KEY = "MDIO__GRID__SPARSITY_RATIO_WARN"
+_GRID_SPARSITY_RATIO_LIMIT_KEY = "MDIO__GRID__SPARSITY_RATIO_LIMIT"
+_SAVE_SEGY_FILE_HEADER_KEY = "MDIO__IMPORT__SAVE_SEGY_FILE_HEADER"
+_MDIO_SEGY_SPEC_KEY = "MDIO__SEGY__SPEC"
+_RAW_HEADERS_KEY = "MDIO__IMPORT__RAW_HEADERS"
+_IGNORE_CHECKS_KEY = "MDIO_IGNORE_CHECKS"
+_CLOUD_NATIVE_KEY = "MDIO__IMPORT__CLOUD_NATIVE"
 
-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)
+# Default values
+_EXPORT_CPUS_DEFAULT = cpu_count(logical=True)
+_IMPORT_CPUS_DEFAULT = cpu_count(logical=True)
+_GRID_SPARSITY_RATIO_WARN_DEFAULT = "2"
+_GRID_SPARSITY_RATIO_LIMIT_DEFAULT = "10"
+_SAVE_SEGY_FILE_HEADER_DEFAULT = "false"
+_MDIO_SEGY_SPEC_DEFAULT = None
+_RAW_HEADERS_DEFAULT = "false"
+_IGNORE_CHECKS_DEFAULT = "false"
+_CLOUD_NATIVE_DEFAULT = "false"
+
+
+def _get_env_value(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)
+
+
+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")
+
+
+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
+
+
+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
+
+
+def export_cpus() -> int:
+    """Number of CPUs to use for export operations."""
+    value = _get_env_value(_EXPORT_CPUS_KEY, _EXPORT_CPUS_DEFAULT)
+    return _parse_int(value, _EXPORT_CPUS_KEY)
+
+
+def import_cpus() -> int:
+    """Number of CPUs to use for import operations."""
+    value = _get_env_value(_IMPORT_CPUS_KEY, _IMPORT_CPUS_DEFAULT)
+    return _parse_int(value, _IMPORT_CPUS_KEY)
+
+
+def grid_sparsity_ratio_warn() -> float:
+    """Sparsity ratio threshold for warnings."""
+    value = _get_env_value(_GRID_SPARSITY_RATIO_WARN_KEY, _GRID_SPARSITY_RATIO_WARN_DEFAULT)
+    return _parse_float(value, _GRID_SPARSITY_RATIO_WARN_KEY)
+
+
+def grid_sparsity_ratio_limit() -> float:
+    """Sparsity ratio threshold for errors."""
+    value = _get_env_value(_GRID_SPARSITY_RATIO_LIMIT_KEY, _GRID_SPARSITY_RATIO_LIMIT_DEFAULT)
+    return _parse_float(value, _GRID_SPARSITY_RATIO_LIMIT_KEY)
+
+
+def save_segy_file_header() -> bool:
+    """Whether to save SEG-Y file headers."""
+    value = _get_env_value(_SAVE_SEGY_FILE_HEADER_KEY, _SAVE_SEGY_FILE_HEADER_DEFAULT)
+    return _parse_bool(value)
+
+
+def mdio_segy_spec() -> str | None:
+    """Path to MDIO SEG-Y specification file."""
+    return _get_env_value(_MDIO_SEGY_SPEC_KEY, _MDIO_SEGY_SPEC_DEFAULT)
+
+
+def raw_headers() -> bool:
+    """Whether to preserve raw headers."""
+    value = _get_env_value(_RAW_HEADERS_KEY, _RAW_HEADERS_DEFAULT)
+    return _parse_bool(value)
+
+
+def ignore_checks() -> bool:
+    """Whether to ignore validation checks."""
+    value = _get_env_value(_IGNORE_CHECKS_KEY, _IGNORE_CHECKS_DEFAULT)
+    return _parse_bool(value)
+
+
+def cloud_native() -> bool:
+    """Whether to use cloud-native mode for SEG-Y processing."""
+    value = _get_env_value(_CLOUD_NATIVE_KEY, _CLOUD_NATIVE_DEFAULT)
+    return _parse_bool(value)

src/mdio/converters/mdio.py

@@ -8,7 +8,7 @@
 import numpy as np
 from tqdm.dask import TqdmCallback
 
-from mdio.api._environ import Environment
+from mdio.api._environ import export_cpus
 from mdio.api.io import _normalize_path
 from mdio.api.io import open_mdio
 from mdio.segy.blocked_io import to_segy
@@ -143,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=Environment.export_cpus())
+                block_records = block_records.compute(num_workers=export_cpus())
 
         ordered_files = [rec.path for rec in block_records.ravel() if rec != 0]
         ordered_files = [output_path] + ordered_files

src/mdio/converters/segy.py

@@ -13,7 +13,11 @@
 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._environ import grid_sparsity_ratio_limit
+from mdio.api._environ import grid_sparsity_ratio_warn
+from mdio.api._environ import ignore_checks as ignore_checks_env
+from mdio.api._environ import raw_headers
+from mdio.api._environ import save_segy_file_header
 from mdio.api.io import _normalize_path
 from mdio.api.io import to_mdio
 from mdio.builder.schemas.chunk_grid import RegularChunkGrid
@@ -99,9 +103,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 = Environment.grid_sparsity_ratio_warn()
-    error_ratio = Environment.grid_sparsity_ratio_limit()
-    ignore_checks = Environment.ignore_checks()
+    warning_ratio = grid_sparsity_ratio_warn()
+    error_ratio = grid_sparsity_ratio_limit()
+    ignore_checks = ignore_checks_env()
 
     # Check sparsity
     should_warn = sparsity_ratio > warning_ratio
@@ -359,7 +363,7 @@ def _populate_coordinates(
 
 
 def _add_segy_file_headers(xr_dataset: xr_Dataset, segy_file_info: SegyFileInfo) -> xr_Dataset:
-    if not Environment.save_segy_file_header():
+    if not save_segy_file_header():
         return xr_dataset
 
     expected_rows = 40
@@ -383,7 +387,7 @@ def _add_segy_file_headers(xr_dataset: xr_Dataset, segy_file_info: SegyFileInfo)
             "binaryHeader": segy_file_info.binary_header_dict,
         }
     )
-    if Environment.raw_headers():
+    if 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})
 
@@ -550,7 +554,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 Environment.raw_headers():
+    if 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:

src/mdio/segy/_workers.py

@@ -8,7 +8,7 @@
 import numpy as np
 from segy.arrays import HeaderArray
 
-from mdio.api._environ import Environment
+from mdio.api._environ import cloud_native
 from mdio.api.io import _normalize_storage_options
 from mdio.segy._raw_trace_wrapper import SegyFileRawTraceWrapper
 from mdio.segy.file import SegyFileArguments
@@ -50,7 +50,7 @@ def header_scan_worker(
 
     slice_ = slice(*trace_range)
 
-    trace_header = segy_file.trace[slice_].header if Environment.cloud_native() else segy_file.header[slice_]
+    trace_header = segy_file.trace[slice_].header if 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.

src/mdio/segy/blocked_io.py

@@ -15,7 +15,7 @@
 from tqdm.auto import tqdm
 from zarr import open_group as zarr_open_group
 
-from mdio.api._environ import Environment
+from mdio.api._environ import import_cpus
 from mdio.api.io import _normalize_storage_options
 from mdio.builder.schemas.v1.stats import CenteredBinHistogram
 from mdio.builder.schemas.v1.stats import SummaryStatistics
@@ -80,7 +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_workers = min(num_chunks, Environment.import_cpus())
+    num_workers = min(num_chunks, import_cpus())
     context = mp.get_context("spawn")
     executor = ProcessPoolExecutor(max_workers=num_workers, mp_context=context)
 

src/mdio/segy/compat.py

@@ -22,7 +22,7 @@
 from segy.schema import TraceSpec
 from segy.standards.fields import binary
 
-from mdio.api._environ import Environment
+from mdio.api._environ import mdio_segy_spec as get_mdio_segy_spec
 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 = Environment.mdio_segy_spec()
+    spec_override = get_mdio_segy_spec()
 
     if spec_override is not None:
         return SegySpec.model_validate_json(spec_override)

src/mdio/segy/parsers.py

@@ -11,7 +11,7 @@
 import numpy as np
 from tqdm.auto import tqdm
 
-from mdio.api._environ import Environment
+from mdio.api._environ import import_cpus
 from mdio.segy._workers import header_scan_worker
 
 if TYPE_CHECKING:
@@ -50,7 +50,7 @@ def parse_headers(
 
         trace_ranges.append((start, stop))
 
-    num_workers = min(n_blocks, Environment.import_cpus())
+    num_workers = min(n_blocks, import_cpus())
 
     tqdm_kw = {"unit": "block", "dynamic_ncols": True}
     # For Unix async writes with s3fs/fsspec & multiprocessing, use 'spawn' instead of default