From f9cfb6ed3bf701bfe8ad5d44e93b102bf9240f0f Mon Sep 17 00:00:00 2001
From: janbjorge <jlov@equinor.com>
Date: Fri, 12 Jan 2024 19:52:09 +0100
Subject: [PATCH] CLN: Remove extension argument (#412)

---
 src/fmu/dataio/_metadata.py |  8 ++++++--
 src/fmu/dataio/_utils.py    | 24 ++++++++++++------------
 src/fmu/dataio/dataio.py    |  3 +--
 3 files changed, 19 insertions(+), 16 deletions(-)

diff --git a/src/fmu/dataio/_metadata.py b/src/fmu/dataio/_metadata.py
index 709f4b3c4..7b9de70c3 100644
--- a/src/fmu/dataio/_metadata.py
+++ b/src/fmu/dataio/_metadata.py
@@ -306,12 +306,16 @@ def _populate_meta_file(self) -> None:
             self.meta_file["absolute_path_symlink"] = fdata.absolute_path_symlink
 
         if self.compute_md5:
-            with NamedTemporaryFile(buffering=0) as tf:
+            if not self.objdata.extension.startswith("."):
+                raise ValueError("A extension must start with '.'")
+            with NamedTemporaryFile(
+                buffering=0,
+                suffix=self.objdata.extension,
+            ) as tf:
                 logger.info("Compute MD5 sum for tmp file...: %s", tf.name)
                 self.meta_file["checksum_md5"] = export_file_compute_checksum_md5(
                     self.obj,
                     Path(tf.name),
-                    self.objdata.extension,
                     flag=self.dataio._usefmtflag,
                 )
         else:
diff --git a/src/fmu/dataio/_utils.py b/src/fmu/dataio/_utils.py
index b83ed2110..a735ccb7f 100644
--- a/src/fmu/dataio/_utils.py
+++ b/src/fmu/dataio/_utils.py
@@ -118,7 +118,6 @@ def export_metadata_file(
 def export_file(
     obj: object,
     filename: Path,
-    extension: str,
     flag: str | None = None,
 ) -> str:
     """Export a valid object to file"""
@@ -126,9 +125,9 @@ def export_file(
     if isinstance(obj, Path):
         # special case when processing data which already has metadata
         shutil.copy(obj, filename)
-    elif extension == ".gri" and isinstance(obj, xtgeo.RegularSurface):
+    elif filename.suffix == ".gri" and isinstance(obj, xtgeo.RegularSurface):
         obj.to_file(filename, fformat="irap_binary")
-    elif extension == ".csv" and isinstance(obj, (xtgeo.Polygons, xtgeo.Points)):
+    elif filename.suffix == ".csv" and isinstance(obj, (xtgeo.Polygons, xtgeo.Points)):
         out = obj.copy()  # to not modify incoming instance!
         assert flag is not None
         if "xtgeo" not in flag:
@@ -139,16 +138,18 @@ def export_file(
                 # out.pname = "ID"  not working
                 out.dataframe.rename(columns={out.pname: "ID"}, inplace=True)
         out.dataframe.to_csv(filename, index=False)
-    elif extension == ".pol" and isinstance(obj, (xtgeo.Polygons, xtgeo.Points)):
+    elif filename.suffix == ".pol" and isinstance(obj, (xtgeo.Polygons, xtgeo.Points)):
         obj.to_file(filename)
-    elif extension == ".segy" and isinstance(obj, xtgeo.Cube):
+    elif filename.suffix == ".segy" and isinstance(obj, xtgeo.Cube):
         obj.to_file(filename, fformat="segy")
-    elif extension == ".roff" and isinstance(obj, (xtgeo.Grid, xtgeo.GridProperty)):
+    elif filename.suffix == ".roff" and isinstance(
+        obj, (xtgeo.Grid, xtgeo.GridProperty)
+    ):
         obj.to_file(filename, fformat="roff")
-    elif extension == ".csv" and isinstance(obj, pd.DataFrame):
+    elif filename.suffix == ".csv" and isinstance(obj, pd.DataFrame):
         includeindex = flag == "include_index"
         obj.to_csv(filename, index=includeindex)
-    elif extension == ".arrow" and HAS_PYARROW and isinstance(obj, pa.Table):
+    elif filename.suffix == ".arrow" and HAS_PYARROW and isinstance(obj, pa.Table):
         # comment taken from equinor/webviz_subsurface/smry2arrow.py
 
         # Writing here is done through the feather import, but could also be done using
@@ -157,11 +158,11 @@ def export_file(
         # actual file format is the same
         # (https://arrow.apache.org/docs/python/feather.html)
         feather.write_feather(obj, dest=filename)
-    elif extension == ".json":
+    elif filename.suffix == ".json":
         with open(filename, "w") as stream:
             json.dump(obj, stream)
     else:
-        raise TypeError(f"Exporting {extension} for {type(obj)} is not supported")
+        raise TypeError(f"Exporting {filename.suffix} for {type(obj)} is not supported")
 
     return str(filename)
 
@@ -177,11 +178,10 @@ def md5sum(fname: Path) -> str:
 def export_file_compute_checksum_md5(
     obj: object,
     filename: Path,
-    extension: str,
     flag: str | None = None,
 ) -> str:
     """Export and compute checksum"""
-    export_file(obj, filename, extension, flag=flag)
+    export_file(obj, filename, flag=flag)
     return md5sum(filename)
 
 
diff --git a/src/fmu/dataio/dataio.py b/src/fmu/dataio/dataio.py
index 888b5970b..59decd690 100644
--- a/src/fmu/dataio/dataio.py
+++ b/src/fmu/dataio/dataio.py
@@ -917,7 +917,6 @@ def export(
         metadata["file"]["checksum_md5"] = export_file_compute_checksum_md5(
             obj,
             outfile,
-            outfile.suffix,
             flag=useflag,  # type: ignore
             # BUG(?): Looks buggy, if flag is bool export_file will blow up.
         )
@@ -1517,7 +1516,7 @@ def export(self, obj: object, **kwargs: object) -> str:
         logger.info("Export to file and compute MD5 sum")
         # inject the computed md5 checksum in metadata
         metadata["file"]["checksum_md5"] = export_file_compute_checksum_md5(
-            obj, outfile, outfile.suffix
+            obj, outfile
         )
 
         export_metadata_file(metafile, metadata, savefmt=self.meta_format)