ENH: Enumerate inplace volumes table columns

src/fmu/dataio/export/rms/_enums.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+from enum import Enum
+from typing import Final
+
+
+class InplaceVolumes:
+    """Enumerations relevant to inplace volumes tables."""
+
+    class Fluid(str, Enum):
+        """Fluid types"""
+
+        OIL = "OIL"
+        GAS = "GAS"
+        WATER = "WATER"
+
+    class TableIndexColumns(str, Enum):
+        """The index columns for an inplace volumes table."""
+
+        FLUID = "FLUID"
+        ZONE = "ZONE"
+        REGION = "REGION"
+        FACIES = "FACIES"
+        LICENSE = "LICENSE"
+
+    FLUID_COLUMN: Final = TableIndexColumns.FLUID
+    """The column name and value used to indicate the index value for fluid type."""
+
+    class VolumetricColumns(str, Enum):
+        """The value columns for an inplace volumes table."""
+
+        BULK = "BULK"
+        NET = "NET"
+        PORV = "PORV"
+        HCPV = "HCPV"
+        STOIIP = "STOIIP"
+        GIIP = "GIIP"
+        ASSOCIATEDGAS = "ASSOCIATEDGAS"
+        ASSOCIATEDOIL = "ASSOCIATEDOIL"
+
+    @staticmethod
+    def index_columns() -> list[str]:
+        """Returns a list of the index columns."""
+        return [k.value for k in InplaceVolumes.TableIndexColumns]
+
+    @staticmethod
+    def value_columns() -> list[str]:
+        """Returns a list of the value columns."""
+        return [k.value for k in InplaceVolumes.VolumetricColumns]
+
+    @staticmethod
+    def table_columns() -> list[str]:
+        """Returns a list of all table columns."""
+        return InplaceVolumes.index_columns() + InplaceVolumes.value_columns()

src/fmu/dataio/export/rms/inplace_volumes.py

@@ -2,7 +2,6 @@
 
 import warnings
 from dataclasses import dataclass
-from enum import Enum
 from pathlib import Path
 from typing import Any, Final
 
@@ -14,6 +13,7 @@
 from fmu.dataio._model.enums import Classification
 from fmu.dataio.export._decorators import experimental
 from fmu.dataio.export._export_result import ExportResult, ExportResultItem
+from fmu.dataio.export.rms import _enums
 from fmu.dataio.export.rms._conditional_rms_imports import import_rms_package
 from fmu.dataio.export.rms._utils import (
     check_rmsapi_version,
@@ -25,27 +25,6 @@
 
 _logger: Final = null_logger(__name__)
 
-_FLUID_COLUMN: Final = "FLUID"
-_TABLE_INDEX_COLUMNS: Final = [_FLUID_COLUMN, "ZONE", "REGION", "FACIES", "LICENSE"]
-_VOLUMETRIC_COLUMNS: Final = [
-    "BULK",
-    "NET",
-    "PORV",
-    "HCPV",
-    "STOIIP",
-    "GIIP",
-    "ASSOCIATEDGAS",
-    "ASSOCIATEDOIL",
-]
-
-
-class _Fluid(str, Enum):
-    """Fluid types"""
-
-    OIL = "OIL"
-    GAS = "GAS"
-    WATER = "WATER"
-
 
 # rename columns to FMU standard
 _RENAME_COLUMNS_FROM_RMS: Final = {
@@ -149,7 +128,7 @@ def _convert_table_from_rms_to_legacy_format(table: pd.DataFrame) -> pd.DataFram
     def _add_missing_columns_to_table(table: pd.DataFrame) -> pd.DataFrame:
         """Add columns with nan values if not present in table."""
         _logger.debug("Add table index columns to table if missing...")
-        for col in _TABLE_INDEX_COLUMNS + _VOLUMETRIC_COLUMNS:
+        for col in _enums.InplaceVolumes.table_columns():
             if col not in table:
                 table[col] = np.nan
         return table
@@ -158,7 +137,7 @@ def _add_missing_columns_to_table(table: pd.DataFrame) -> pd.DataFrame:
     def _set_table_column_order(table: pd.DataFrame) -> pd.DataFrame:
         """Set the column order in the table."""
         _logger.debug("Settting the table column order...")
-        return table[_TABLE_INDEX_COLUMNS + _VOLUMETRIC_COLUMNS]
+        return table[_enums.InplaceVolumes.table_columns()]
 
     @staticmethod
     def _transform_and_add_fluid_column_to_table(
@@ -170,10 +149,15 @@ def _transform_and_add_fluid_column_to_table(
         are renamed into 'BULK' and 'PORV' columns. To separate the data an additional
         FLUID column is added that indicates the type of fluid the row represents.
         """
-        table_index = [col for col in _TABLE_INDEX_COLUMNS if col in table]
+        table_index = [
+            col for col in _enums.InplaceVolumes.index_columns() if col in table
+        ]
 
         tables = []
-        for fluid in [_Fluid.GAS.value, _Fluid.OIL.value]:
+        for fluid in [
+            _enums.InplaceVolumes.Fluid.GAS.value,
+            _enums.InplaceVolumes.Fluid.OIL.value,
+        ]:
             fluid_columns = [col for col in table.columns if col.endswith(f"_{fluid}")]
             if fluid_columns:
                 fluid_table = table[table_index + fluid_columns].copy()
@@ -182,7 +166,7 @@ def _transform_and_add_fluid_column_to_table(
                 fluid_table.columns = fluid_table.columns.str.replace(f"_{fluid}", "")
 
                 # add the fluid as column entry instead
-                fluid_table[_FLUID_COLUMN] = fluid.lower()
+                fluid_table[_enums.InplaceVolumes.FLUID_COLUMN] = fluid.lower()
 
                 tables.append(fluid_table)
 
@@ -196,7 +180,9 @@ def _convert_table_from_legacy_to_standard_format(
         product. The standard format has a fluid column, and all table_index and
         volumetric columns are present with a standard order in the table.
         """
-        table_index = [col for col in _TABLE_INDEX_COLUMNS if col in table]
+        table_index = [
+            col for col in _enums.InplaceVolumes.index_columns() if col in table
+        ]
         table = self._transform_and_add_fluid_column_to_table(table, table_index)
         table = self._add_missing_columns_to_table(table)
         return self._set_table_column_order(table)
@@ -212,8 +198,8 @@ def _validate_table(self) -> None:
         """
         _logger.debug("Validating the dataframe...")
 
-        has_oil = "oil" in self._dataframe[_FLUID_COLUMN].values
-        has_gas = "gas" in self._dataframe[_FLUID_COLUMN].values
+        has_oil = "oil" in self._dataframe[_enums.InplaceVolumes.FLUID_COLUMN].values
+        has_gas = "gas" in self._dataframe[_enums.InplaceVolumes.FLUID_COLUMN].values
 
         # check that one of oil and gas fluids are present
         if not (has_oil or has_gas):
@@ -255,7 +241,7 @@ def _export_volume_table(self) -> ExportResult:
             classification=self._classification,
             name=self.grid_name,
             rep_include=False,
-            table_index=_TABLE_INDEX_COLUMNS,
+            table_index=_enums.InplaceVolumes.index_columns(),
         )
         absolute_export_path = edata.export(self._dataframe)
         _logger.debug("Volume result to: %s", absolute_export_path)

tests/test_export_rms/test_export_rms_volumetrics.py

@@ -94,13 +94,13 @@ def test_rms_volumetrics_export_class(exportvolumetrics):
 def test_rms_volumetrics_export_class_table_index(voltable_standard, exportvolumetrics):
     """See mocks in local conftest.py"""
 
-    from fmu.dataio.export.rms.inplace_volumes import _TABLE_INDEX_COLUMNS
+    from fmu.dataio.export.rms import _enums
 
     out = exportvolumetrics._export_volume_table()
     metadata = dataio.read_metadata(out.items[0].absolute_path)
 
     # check that the table index is set correctly
-    assert metadata["data"]["table_index"] == _TABLE_INDEX_COLUMNS
+    assert metadata["data"]["table_index"] == _enums.InplaceVolumes.index_columns()
 
     # should fail if missing table index
     exportvolumetrics._dataframe = voltable_standard.drop(columns="ZONE")
@@ -120,10 +120,8 @@ def test_convert_table_from_legacy_to_standard_format(
     """Test that a voltable with legacy format is converted to
     the expected standard format"""
 
-    from fmu.dataio.export.rms.inplace_volumes import (
-        _FLUID_COLUMN,
-        _ExportVolumetricsRMS,
-    )
+    from fmu.dataio.export.rms import _enums
+    from fmu.dataio.export.rms.inplace_volumes import _ExportVolumetricsRMS
 
     monkeypatch.chdir(rmssetup_with_fmuconfig)
 
@@ -147,8 +145,11 @@ def test_convert_table_from_legacy_to_standard_format(
     pd.testing.assert_frame_equal(voltable_standard, exported_table)
 
     # check that the fluid column exists and contains oil and gas
-    assert _FLUID_COLUMN in exported_table
-    assert set(exported_table[_FLUID_COLUMN].unique()) == {"oil", "gas"}
+    assert _enums.InplaceVolumes.FLUID_COLUMN in exported_table
+    assert set(exported_table[_enums.InplaceVolumes.FLUID_COLUMN].unique()) == {
+        "oil",
+        "gas",
+    }
 
     # check the column order
     assert list(exported_table.columns) == EXPECTED_COLUMN_ORDER
@@ -298,11 +299,8 @@ def test_rms_volumetrics_export_function(
 ):
     """Test the public function."""
 
-    from fmu.dataio.export.rms import export_inplace_volumes
-    from fmu.dataio.export.rms.inplace_volumes import (
-        _TABLE_INDEX_COLUMNS,
-        _ExportVolumetricsRMS,
-    )
+    from fmu.dataio.export.rms import _enums, export_inplace_volumes
+    from fmu.dataio.export.rms.inplace_volumes import _ExportVolumetricsRMS
 
     monkeypatch.chdir(rmssetup_with_fmuconfig)
 
@@ -327,7 +325,7 @@ def test_rms_volumetrics_export_function(
 
     assert "volumes" in metadata["data"]["content"]
     assert metadata["access"]["classification"] == "restricted"
-    assert metadata["data"]["table_index"] == _TABLE_INDEX_COLUMNS
+    assert metadata["data"]["table_index"] == _enums.InplaceVolumes.index_columns()
 
 
 @inside_rms
@@ -360,14 +358,10 @@ def test_inplace_volumes_payload_validates_against_schema(
 
 @inside_rms
 def test_inplace_volumes_export_and_result_columns_are_the_same(
-    mock_project_variable,
     mocked_rmsapi_modules,
 ) -> None:
-    from fmu.dataio.export.rms.inplace_volumes import (
-        _TABLE_INDEX_COLUMNS,
-        _VOLUMETRIC_COLUMNS,
-    )
+    from fmu.dataio.export.rms import _enums
 
-    export_columns = _TABLE_INDEX_COLUMNS + _VOLUMETRIC_COLUMNS
-    result_columns = InplaceVolumesResultRow.model_fields.keys()
-    assert set(export_columns) == set(result_columns)
+    assert _enums.InplaceVolumes.table_columns() == list(
+        InplaceVolumesResultRow.model_fields.keys()
+    )