refactor: pickle utils

queens/utils/cli.py

@@ -23,10 +23,10 @@
 from queens.utils.exceptions import CLIError
 from queens.utils.injector import inject
 from queens.utils.input_to_script import create_script_from_input_file
+from queens.utils.io import print_pickled_data
 from queens.utils.logger_settings import reset_logging, setup_cli_logging
 from queens.utils.metadata import write_metadata_to_csv
 from queens.utils.path import PATH_TO_QUEENS
-from queens.utils.pickle import print_pickled_data
 from queens.utils.printing import get_str_table
 from queens.utils.run_subprocess import run_subprocess
 

queens/utils/io.py

@@ -15,19 +15,87 @@
 """Utils for input/output handling."""
 
 import csv
+import logging
+import pickle
 from pathlib import Path
 
 import yaml
 
 from queens.utils.exceptions import FileTypeError
-from queens.utils.pickle import load_pickle
 
 try:
     import simplejson as json
 except ImportError:
     import json
 
 
+_logger = logging.getLogger(__name__)
+
+
+def load_pickle(file_path):
+    """Load a pickle file directly from path.
+
+    Args:
+        file_path (Path): Path to pickle-file
+
+    Returns:
+        Data (dict) in the pickle file
+    """
+    if not file_path.is_file():
+        raise FileNotFoundError(f"File {file_path} does not exist.")
+    try:
+        data = pickle.load(file_path.open("rb"))
+        return data
+    except Exception as exception:
+        raise IOError(f"Could not open the pickle file {file_path}") from exception
+
+
+def print_pickled_data(file_path):
+    """Print a table of the data within a pickle file.
+
+    Only goes one layer deep for dicts. This is similar to *python -m pickle file_path* but makes
+    it a single command and fancy prints.
+
+    Args:
+        file_path (Path): Path to pickle-file
+    """
+    data = load_pickle(file_path)
+    _logger.info("\n\npickle file: %s", file_path)
+    for key, item in data.items():
+        item_type = type(item)
+        if isinstance(item, dict):
+            string = ""
+            for subkey, subitem in item.items():
+                string += (
+                    _create_single_item_string(subkey, subitem, type(subitem), seperator="-") + "\n"
+                )
+            item = string.replace("\n", "\n    ")
+        _logger.info(_create_single_item_string(key, item, item_type))
+        _logger.info(" ")
+
+
+def _create_single_item_string(key, item, item_type, seperator="="):
+    """Create a table for a single item.
+
+    Args:
+        key (str): Key of the item
+        item (obj): Item value for the key
+        item_type (str): Type of the item value
+        seperator (str, optional): Create seperator line (default is "=")
+
+    Returns:
+        string: table for this item.
+    """
+    string = (
+        seperator * 60
+        + f"\nKey:  {key}\n"
+        + f"Type: {item_type}\n"
+        + f"Value:\n{item}\n"
+        + seperator * 60
+    )
+    return string
+
+
 def load_input_file(input_file_path):
     """Load inputs from file by path.