Refactoring

Signed-off-by: zethson <[email protected]>

ehrapy/io/_write.py

@@ -6,12 +6,13 @@
 import numpy as np
 
 from ehrapy import settings
-from ehrapy.io._utility_io import _get_file_extension
 from ehrapy.preprocessing._encode import encode
 
 if TYPE_CHECKING:
     from anndata import AnnData
 
+supported_extensions = {"csv", "tsv", "h5ad"}
+
 
 def write(
     filename: str | Path,
@@ -69,6 +70,31 @@ def write(
             adata.write(filename, compression=compression, compression_opts=compression_opts)
 
 
+def _get_file_extension(file_path: Path) -> str:
+    """Check whether the argument is a filename.
+
+    Args:
+        file_path: Path to the file.
+
+    Returns:
+        File extension of the specified file
+    """
+    ext = file_path.suffixes
+
+    if len(ext) > 2:
+        ext = ext[-2:]
+
+    if ext and ext[-1][1:] in supported_extensions:
+        return ext[-1][1:]
+    raise ValueError(
+        f"""\
+        {file_path!r} does not end on a valid extension.
+        Please, provide one of the available extensions.
+        {supported_extensions}
+        """
+    )
+
+
 def _get_filename_from_key(key, extension=None) -> Path:
     """Gets full file name from a key.