From 7e8d4e00c8baf645e8984443fa00ae9ab69377ac Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Therese=20Natter=C3=B8y?=
 <61694854+tnatt@users.noreply.github.com>
Date: Thu, 19 Dec 2024 13:41:15 +0100
Subject: [PATCH] MAINT: Replace column strings with enums in inplace_volumes
 (#936)

---
 src/fmu/dataio/export/rms/inplace_volumes.py | 67 ++++++++++++--------
 1 file changed, 41 insertions(+), 26 deletions(-)

diff --git a/src/fmu/dataio/export/rms/inplace_volumes.py b/src/fmu/dataio/export/rms/inplace_volumes.py
index 66436cacc..60ad4619d 100644
--- a/src/fmu/dataio/export/rms/inplace_volumes.py
+++ b/src/fmu/dataio/export/rms/inplace_volumes.py
@@ -27,28 +27,31 @@
 _logger: Final = null_logger(__name__)
 
 
+_VolumetricColumns = _enums.InplaceVolumes.VolumetricColumns
+_TableIndexColumns = _enums.InplaceVolumes.TableIndexColumns
+
 # rename columns to FMU standard
 _RENAME_COLUMNS_FROM_RMS: Final = {
     "Proj. real.": "REAL",
-    "Zone": "ZONE",
-    "Segment": "REGION",
-    "Boundary": "LICENSE",
-    "Facies": "FACIES",
-    "BulkOil": "BULK_OIL",
-    "NetOil": "NET_OIL",
-    "PoreOil": "PORV_OIL",
-    "HCPVOil": "HCPV_OIL",
-    "STOIIP": "STOIIP_OIL",
-    "AssociatedGas": "ASSOCIATEDGAS_OIL",
-    "BulkGas": "BULK_GAS",
-    "NetGas": "NET_GAS",
-    "PoreGas": "PORV_GAS",
-    "HCPVGas": "HCPV_GAS",
-    "GIIP": "GIIP_GAS",
-    "AssociatedLiquid": "ASSOCIATEDOIL_GAS",
-    "Bulk": "BULK_TOTAL",
-    "Net": "NET_TOTAL",
-    "Pore": "PORV_TOTAL",
+    "Zone": _TableIndexColumns.ZONE.value,
+    "Segment": _TableIndexColumns.REGION.value,
+    "Boundary": _TableIndexColumns.LICENSE.value,
+    "Facies": _TableIndexColumns.FACIES.value,
+    "BulkOil": _VolumetricColumns.BULK.value + "_OIL",
+    "NetOil": _VolumetricColumns.NET.value + "_OIL",
+    "PoreOil": _VolumetricColumns.PORV.value + "_OIL",
+    "HCPVOil": _VolumetricColumns.HCPV.value + "_OIL",
+    "STOIIP": _VolumetricColumns.STOIIP.value + "_OIL",
+    "AssociatedGas": _VolumetricColumns.ASSOCIATEDGAS.value + "_OIL",
+    "BulkGas": _VolumetricColumns.BULK.value + "_GAS",
+    "NetGas": _VolumetricColumns.NET.value + "_GAS",
+    "PoreGas": _VolumetricColumns.PORV.value + "_GAS",
+    "HCPVGas": _VolumetricColumns.HCPV.value + "_GAS",
+    "GIIP": _VolumetricColumns.GIIP.value + "_GAS",
+    "AssociatedLiquid": _VolumetricColumns.ASSOCIATEDOIL.value + "_GAS",
+    "Bulk": _VolumetricColumns.BULK.value + "_TOTAL",
+    "Net": _VolumetricColumns.NET.value + "_TOTAL",
+    "Pore": _VolumetricColumns.PORV.value + "_TOTAL",
 }
 
 
@@ -240,8 +243,14 @@ def _validate_table(self) -> None:
         """
         _logger.debug("Validating the dataframe...")
 
-        has_oil = "oil" in self._dataframe[_enums.InplaceVolumes.FLUID_COLUMN].values
-        has_gas = "gas" in self._dataframe[_enums.InplaceVolumes.FLUID_COLUMN].values
+        has_oil = (
+            _enums.InplaceVolumes.Fluid.oil.value
+            in self._dataframe[_enums.InplaceVolumes.FLUID_COLUMN].values
+        )
+        has_gas = (
+            _enums.InplaceVolumes.Fluid.gas.value
+            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):
@@ -253,15 +262,21 @@ def _validate_table(self) -> None:
 
         # create list of missing or non-defined required columns
         missing_calculations = []
-        for col in ["BULK", "PORV", "HCPV"]:
+        for col in [
+            _VolumetricColumns.BULK.value,
+            _VolumetricColumns.PORV.value,
+            _VolumetricColumns.HCPV.value,
+        ]:
             if self._is_column_missing_in_table(col):
                 missing_calculations.append(col)
 
-        if has_oil and self._is_column_missing_in_table("STOIIP"):
-            missing_calculations.append("STOIIP")
+        if has_oil and self._is_column_missing_in_table(
+            _VolumetricColumns.STOIIP.value
+        ):
+            missing_calculations.append(_VolumetricColumns.STOIIP.value)
 
-        if has_gas and self._is_column_missing_in_table("GIIP"):
-            missing_calculations.append("GIIP")
+        if has_gas and self._is_column_missing_in_table(_VolumetricColumns.GIIP.value):
+            missing_calculations.append(_VolumetricColumns.GIIP.value)
 
         if missing_calculations:
             raise RuntimeError(