mfile cleanup

Author: je-cook

process/core/io/mfile/cli.py

@@ -1,8 +1,10 @@
+import json
+
 import click
 
+from process.core.io.mfile import MFile
 from process.core.io.mfile.mfile_comparison import compare_mfiles
-from process.core.io.mfile.mfile_to_csv import to_csv
-from process.core.io.tools import mfile_arg, mfile_opt, save
+from process.core.io.tools import mfile_arg, mfile_opt, save, scan_opt
 
 
 @click.group()
@@ -16,17 +18,27 @@ def mfile():
     "-v",
     "--variables",
     type=str,
-    help="Optional list of variables or file with list of variables to extract",
+    help="Optional list of variables or json file with list of variables to extract",
 )
 @click.option(
     "-fmt",
     "--format",
     "format_",
-    type=click.Choice(["json", "csv", "toml", "yaml", "pickle"]),
+    type=click.Choice(["json", "csv", "toml"]),
 )
-def convert(mfile, variables, format_):
+@scan_opt
+@click.option("--verbose", is_flag=True)
+def convert(mfile, variables, format_, scan, verbose):
     """Convert MFile to other formats."""
-    to_csv(mfile, variables)
+    if variables.endswith(".json"):
+        with open(variables) as f:
+            variables = json.load(f)["variables"]
+    else:
+        variables = list(filter(None, variables.replace(" ", ":").split(":")))
+
+    getattr(MFile(mfile), f"to_{format_}")(
+        keys_to_write=variables, scan=scan, verbose=verbose
+    )
 
 
 @mfile.command("compare", no_args_is_help=True)

process/core/io/mfile/mfile.py

@@ -27,6 +27,8 @@
 from collections import OrderedDict
 from typing import Any
 
+import numpy as np
+
 logger = logging.getLogger(__name__)
 
 
@@ -337,52 +339,118 @@ def add_to_mfile_variable(self, des, name, value, unit, flag, scan=None):
             self.data[var_key] = var
             self.data[var_key].set_scan(1, value)
 
-    def write_to_json(self, keys_to_write=None, scan=-1, verbose=False):
+    def to_dict(self, keys=None, scan: int | None = -1, verbose=False) -> dict:
+        """Convert MFile to dictionary
+
+        Parameters
+        ----------
+        keys :
+             keys to select
+        scan :
+             scan to select
+        verbose :
+             verbosity of output
+        """
+
+        if keys is None:
+            keys = self.data.keys()
+
+        def _get_data(item, dat_key):
+            data = self.data[item].get_scan(dat_key)
+            des = self.data[item].var_description.replace("_", " ")
+            return {"value": data, "description": des} if verbose else data
+
+        save_range = (
+            range(1, self.data["rmajor"].get_number_of_scans() + 1)
+            if scan is None
+            else [scan]
+        )
+        output = {
+            f"scan-{i + 1}": {
+                item: _get_data(
+                    item, -1 if self.data[item].get_number_of_scans() == 1 else i
+                )
+                for item in keys
+            }
+            for i in save_range
+        }
+        return (
+            output[next(iter(output.keys()))]
+            if len(output.keys()) == 1 and scan is not None
+            else output
+        )
+
+    def to_json(self, keys_to_write=None, scan: int | None = -1, verbose=False):
         """Write MFILE object to JSON file
 
         Parameters
         ----------
         keys_to_write :
-             (Default value = None)
+             keys to select
         scan :
-             (Default value = -1)
+             scan to select
         verbose :
-             (Default value = False)
+             verbosity of output
         """
 
-        if keys_to_write is None:
-            keys_to_write = self.data.keys()
-
         filename = f"{self.filename}.json"
 
-        dict_to_write = {}
-
-        if scan == 0:
-            for i in range(self.data["rmajor"].get_number_of_scans()):
-                sub_dict = {}
-                for item in keys_to_write:
-                    dat_key = -1 if self.data[item].get_number_of_scans() == 1 else i + 1
-                    data = self.data[item].get_scan(dat_key)
-                    des = self.data[item].var_description.replace("_", " ")
-                    entry = {"value": data, "description": des} if verbose else data
-                    sub_dict[item] = entry
-                dict_to_write[f"scan-{i + 1}"] = sub_dict
-        else:
-            for item in keys_to_write:
-                # Initialize dat_key properly based on the number of scans
-                if self.data[item].get_number_of_scans() == 1:
-                    dat_key = -1
-                else:
-                    dat_key = (
-                        scan if scan > 0 else 1
-                    )  # Default to scan 1 if not specified
-                data = self.data[item].get_scan(dat_key)
-                des = self.data[item].var_description.replace("_", " ")
-                entry = {"value": data, "description": des} if verbose else data
-                dict_to_write[item] = entry
-
         with open(filename, "w") as fp:
-            json.dump(dict_to_write, fp, indent=4)
+            json.dump(self.to_dict(keys_to_write, scan, verbose), fp, indent=4)
+
+    def to_toml(self, keys_to_write=None, scan: int | None = -1, verbose=False):
+        """Write MFILE object to JSON file
+
+        Parameters
+        ----------
+        keys_to_write :
+             keys to select
+        scan :
+             scan to select
+        verbose :
+             verbosity of output
+        """
+        import toml
+
+        with open(f"{self.filename}.toml", "w") as file:
+            toml.dump(self.to_dict(keys_to_write, scan, verbose), file)
+
+    def to_csv(self, keys_to_write=None, scan=-1, verbose=False):
+        """Write to csv file.
+
+        Parameters
+        ----------
+        args : string, list of tuples
+            input filename, variable data
+        csv_outfile :
+
+        output_data :
+             (Default value = None)
+        """
+        output_data = []
+        if scan is None:
+            for scan_key, vals in self.to_dict(
+                keys_to_write, scan=scan, verbose=verbose
+            ).items():
+                output_data.extend((
+                    (scan_key, "", ""),
+                    ("Description", "Varname", "Value"),
+                ))
+                for k, v in vals.items():
+                    output_data.append((v["description"], k, v["value"]))
+        else:
+            output_data.append(("Description", "Varname", "Value"))
+            for k, v in self.to_dict(keys_to_write, scan=scan, verbose=verbose).items():
+                output_data.append((v["description"], k, v["value"]))
+        np.savetxt(
+            f"{self.filename}.csv",
+            output_data or [],
+            fmt="%.5e",
+            delimiter=",",
+            header="PROCESS MFILE converted to csv",
+            footer="",
+            comments="",
+        )
 
 
 def sort_value(value_words: list[str]) -> str | float: