MAINT: update file extension logic (bilby-dev#962)

* MAINT: rework file extension logic

* TST: update result saving tests

* BUG: fix extension precedence in `read_in_result`

* TST: assert files can be read in

* MAINT: tweak extension logic in `read_in_result`

* TST: update extensions in tests

* BUG: fix extension default

* BUG: fix logic in `read_in_result`

* TST: improve result save/read tests

* ENH: improve error handling in `read_in_result`

* TST: add tests for file loading with incorrect extensions

* MAINT: tweak IOError message

* TST: fix logging check in caching test

---------

Co-authored-by: Colm Talbot <talbotcolm@gmail.com>

Author: mj-will

bilby/core/result.py

@@ -82,7 +82,7 @@ def _determine_file_name(filename, outdir, label, extension, gzip):
             return result_file_name(outdir, label, extension, gzip)
 
 
-def read_in_result(filename=None, outdir=None, label=None, extension='json', gzip=False, result_class=None):
+def read_in_result(filename=None, outdir=None, label=None, extension=None, gzip=False, result_class=None):
     """ Reads in a stored bilby result object
 
     Parameters
@@ -92,33 +92,67 @@ def read_in_result(filename=None, outdir=None, label=None, extension='json', gzi
     outdir, label, extension: str
         Name of the output directory, label and extension used for the default
         naming scheme.
+    extension: str, optional
+        The file extension to use. If not given, the extension is inferred from
+        the filename if provided. If the filename is not given, defaults to
+        'json'.
     result_class: bilby.core.result.Result, or child of
         The result class to use. By default, `bilby.core.result.Result` is used,
         but objects which inherit from this class can be given providing
         additional methods.
     """
-    filename = _determine_file_name(filename, outdir, label, extension, gzip)
+    filename = _determine_file_name(
+        filename, outdir, label, extension or "json", gzip
+    )
 
     if result_class is None:
         result_class = Result
     elif not issubclass(result_class, Result):
         raise ValueError(f"Input result_class={result_class} not understood")
 
     # Get the actual extension (may differ from the default extension if the filename is given)
-    extension = os.path.splitext(filename)[1].lstrip('.')
+    if extension is None:
+        ext = os.path.splitext(filename)[1][1:]
+        extension = ext if ext else None
     if extension == 'gz':  # gzipped file
         extension = os.path.splitext(os.path.splitext(filename)[0])[1].lstrip('.')
 
-    if 'json' in extension:
-        result = result_class.from_json(filename=filename)
-    elif ('hdf5' in extension) or ('h5' in extension):
-        result = result_class.from_hdf5(filename=filename)
-    elif ("pkl" in extension) or ("pickle" in extension):
-        result = result_class.from_pickle(filename=filename)
-    elif extension is None:
-        raise ValueError("No filetype extension provided")
-    else:
-        raise ValueError("Filetype {} not understood".format(extension))
+    if extension is None:
+        raise ValueError("No filetype extension provided and could not be inferred from filename")
+
+    extension = extension.lower()
+    read_functions = {
+        'json': result_class.from_json,
+        'hdf5': result_class.from_hdf5,
+        'h5': result_class.from_hdf5,
+        'pkl': result_class.from_pickle,
+        'pickle': result_class.from_pickle,
+    }
+    if extension not in read_functions:
+        raise ValueError(
+            f"Filetype {extension} not understood, known types are {list(read_functions.keys())}"
+        )
+
+    func = read_functions[extension]
+
+    # Raise IO errors since this is what the caching relies on
+    # Catch all other exceptions and raise a FileLoadError
+    try:
+        result = func(filename=filename)
+    except IOError as e:
+        raise IOError(
+            f"Failed to read in file {filename} using "
+            f"`{result_class.__name__}.{func.__name__}` "
+            f"(extension={extension}). "
+            "This is likely because the file does not exist or is not a valid bilby result."
+            f"The error was: {e}"
+        ) from e
+    except Exception as e:
+        raise FileLoadError(
+            f"Failed to read in file {filename} using "
+            f"`{result_class.__name__}.{func.__name__}` "
+            f"(extension={extension}). The error was: {e}"
+        ) from e
     return result
 
 
@@ -777,7 +811,8 @@ def save_to_file(self, filename=None, overwrite=False, outdir=None,
         Parameters
         ==========
         filename: optional,
-            Filename to write to (overwrites the default)
+            Filename to write to (overwrites the default). Assumed to include the
+            file extension.
         overwrite: bool, optional
             Whether or not to overwrite an existing result file.
             default=False
@@ -794,24 +829,39 @@ def save_to_file(self, filename=None, overwrite=False, outdir=None,
 
         _outdir = None
         if filename is not None:
-            _outdir, base_filename = os.path.split(filename)
+            # Overwrite the full filename with the base filename
+            # Well append the outdir later on
+            _outdir, filename = os.path.split(filename)
             _outdir = None if _outdir == "" else _outdir
-            base, ext = os.path.splitext(base_filename)
-            if extension is None:
-                extension = ext[1:] if ext else "json"
-                filename = base_filename
-            elif extension is True:
-                message = "Result.save_to_file called with extension=True. "
-                if len(ext) > 0:
-                    message += f"Overwriting extension to json from {ext}, this"
-                else:
-                    message += "This"
-                message += " behaviour is deprecated and will be removed."
-                logger.warning(message)
-                extension = "json"
-                filename = base_filename
-            filename = f"{base}.{extension}"
-        if extension is None or extension is True:
+            _, ext = os.path.splitext(filename)
+            ext = ext[1:] if ext else None
+            # If the extension has not been set, try to infer it from the filename
+            # if not, it will fall back to the default
+            if ext in EXTENSIONS:
+                if extension is None:
+                    logger.debug(
+                        f"Inferred extension '{ext}' from filename '{filename}'. "
+                        "Using this extension for saving."
+                    )
+                    extension = ext
+                elif ext != extension:
+                    message = (
+                        f"The specified extension '{ext}' "
+                        f"does not match the provided extension '{extension}'. "
+                    )
+                    logger.warning(message)
+
+        if extension is None:
+            logger.info("No extension given, defaulting to JSON.")
+            extension = 'json'
+
+        if extension is True:
+            message = (
+                "Result.save_to_file called with extension=True. "
+                "This will default to json, and ignore the extension from the filename. "
+                "This behaviour is deprecated and will be removed. "
+            )
+            logger.warning(message)
             extension = 'json'
 
         outdir = _outdir if outdir is None else outdir
@@ -845,13 +895,13 @@ def save_to_file(self, filename=None, overwrite=False, outdir=None,
                 else:
                     with open(output_path, 'w') as file:
                         json.dump(dictionary, file, indent=2, cls=BilbyJsonEncoder)
-            elif extension == 'hdf5':
+            elif extension in ['hdf5', 'h5']:
                 import h5py
                 dictionary["__module__"] = self.__module__
                 dictionary["__name__"] = self.__class__.__name__
                 with h5py.File(output_path, 'w') as h5file:
                     recursively_save_dict_contents_to_group(h5file, '/', dictionary)
-            elif extension == 'pkl':
+            elif extension in ['pkl', 'pickle']:
                 safe_file_dump(self, output_path, "dill")
             else:
                 raise ValueError(f"Extension type {extension} not understood")
@@ -2291,3 +2341,7 @@ class ResultListError(ResultError):
 
 class FileMovedError(ResultError):
     """ Exceptions that occur when files have been moved """
+
+
+class FileLoadError(ResultError):
+    """ Exceptions that occur when files cannot be loaded """