fixup! Support parquet in sim2seis mock

Author: akustn

test-data/ert/seismic_example/mock_sim2seis.py

@@ -54,6 +54,11 @@ class MonitorDate(StrEnum):
     JAN2030 = "20300101"
 
 
+class OutputFormat(StrEnum):
+    CSV = "csv"
+    PARQUET = "parquet"
+
+
 # Coordinate/region table: (X_UTME, Y_UTMN, REGION)
 COORDINATE_TABLE: list[tuple[float, float, float]] = [
     (463401.665023891, 6929758.90312445, 1.0),
@@ -177,9 +182,9 @@ def _build_stem(
 
 
 def generate_file(
-    output_format: str,
     parameters: dict[str, dict[str, float]],
     output_dir: Path,
+    output_format: OutputFormat,
     horizon: HorizonName,
     attribute: Attribute,
     stacking_offset: StackingOffset,
@@ -218,7 +223,7 @@ def generate_file(
         )
 
     df = pd.DataFrame(rows)
-    if output_format == "parquet":
+    if output_format is OutputFormat.PARQUET:
         df.to_parquet(filepath, index=False)
     else:
         df.to_csv(filepath, index=False)
@@ -229,7 +234,7 @@ def generate_file(
 def generate_many(
     parameters: dict[str, dict[str, float]],
     output_dir: Path,
-    output_format: str,
+    output_format: OutputFormat,
     horizons: list[HorizonName] | None = None,
     attributes: list[Attribute] | None = None,
     stacking_offsets: list[StackingOffset] | None = None,
@@ -259,9 +264,9 @@ def generate_many(
         monitors,
     ):
         generate_file(
-            output_format,
             parameters,
             output_dir,
+            output_format,
             *combo,
         )
 
@@ -280,8 +285,9 @@ def generate_many(
     )
     parser.add_argument(
         "--format",
-        choices=["csv", "parquet"],
-        default="csv",
+        type=OutputFormat,
+        choices=list(OutputFormat),
+        default=OutputFormat.CSV,
         help="Output file format.",
     )
     args = parser.parse_args()