disdrodb 0.0.21__py3-none-any.whl → 0.1.1__py3-none-any.whl

disdrodb/utils/directories.py

@@ -23,34 +23,107 @@ import logging
 import os
 import pathlib
 import shutil
+from typing import Union
+
+from disdrodb.utils.list import flatten_list
+from disdrodb.utils.logger import log_info
 
 logger = logging.getLogger(__name__)
 
 
 def ensure_string_path(path, msg, accepth_pathlib=False):
-    if accepth_pathlib:
-        valid_types = (str, pathlib.PurePath)
-    else:
-        valid_types = str
+    """Ensure that the path is a string."""
+    valid_types = (str, pathlib.PurePath) if accepth_pathlib else str
     if not isinstance(path, valid_types):
         raise TypeError(msg)
     return str(path)
 
 
+def contains_netcdf_or_parquet_files(dir_path: str) -> bool:
+    """Check (recursively) if a directory has any Parquet or netCDF file.
+
+    os.walk under the hood uses os.scandir
+    os.walk file generator + any() avoid use of while loop
+
+    The function returns True as soon as one file is found (short-circuit)^; False otherwise.
+    """
+    suffixes = (".nc", ".parquet")
+    return any(fname.endswith(suffixes) for _, _, files in os.walk(dir_path) for fname in files)
+
+
+def contains_files(dir_path: str) -> bool:
+    """Check (recursively) if a directory contains any file.
+
+    os.walk under the hood uses os.scandir
+    os.walk file generator + any() avoid use of while loop
+
+    The function returns True as soon as one file is found (short-circuit); False otherwise.
+    """
+    return any(fname for _, _, files in os.walk(dir_path) for fname in files)
+
+
+def check_glob_pattern(pattern: str) -> None:
+    """Check if glob pattern is a string and is a valid pattern.
+
+    Parameters
+    ----------
+    pattern : str
+        String to be checked.
+    """
+    if not isinstance(pattern, str):
+        raise TypeError("Expect pattern as a string.")
+    if pattern[0] == "/":
+        raise ValueError("glob_pattern should not start with /")
+    if "//" in pattern:
+        raise ValueError("glob_pattern expects path with single separators: /, not //")
+    if "\\" in pattern:
+        raise ValueError("glob_pattern expects path separators to be /, not \\")
+    return pattern
+
+
+def check_glob_patterns(patterns: Union[str, list]) -> list:
+    """Check if glob patterns are valids."""
+    if not isinstance(patterns, (str, list)):
+        raise ValueError("'glob_patterns' must be a str or list of strings.")
+    if isinstance(patterns, str):
+        patterns = [patterns]
+    patterns = [check_glob_pattern(pattern) for pattern in patterns]
+    return patterns
+
+
 def _recursive_glob(dir_path, glob_pattern):
-    # ** search for all files recursively
-    # glob_pattern = os.path.join(base_dir, "**", "metadata", f"{station_name}.yml")
-    # metadata_filepaths = glob.glob(glob_pattern, recursive=True)
+    # ** search for in zero or all subdirectories recursively
 
     dir_path = pathlib.Path(dir_path)
     return [str(path) for path in dir_path.rglob(glob_pattern)]
 
 
-def list_paths(dir_path, glob_pattern, recursive=False):
+def _list_paths(dir_path, glob_pattern, recursive=False):
+    """Return a list of filepaths and directory paths based on a single glob pattern."""
+    # If glob pattern has separators, disable recursive option
+    if "/" in glob_pattern and "**" not in glob_pattern:
+        recursive = False
+    # Search paths
     if not recursive:
         return glob.glob(os.path.join(dir_path, glob_pattern))
-    else:
-        return _recursive_glob(dir_path, glob_pattern)
+    return _recursive_glob(dir_path, glob_pattern)
+
+
+def list_paths(dir_path, glob_pattern, recursive=False):
+    """Return a list of filepaths and directory paths.
+
+    This function accept also a list of glob patterns !
+    """
+    # Check validity of glob pattern(s)
+    glob_patterns = check_glob_patterns(glob_pattern)
+    # Search path for specified glob patterns
+    paths = flatten_list(
+        [
+            _list_paths(dir_path=dir_path, glob_pattern=glob_pattern, recursive=recursive)
+            for glob_pattern in glob_patterns
+        ],
+    )
+    return paths
 
 
 def list_files(dir_path, glob_pattern, recursive=False):
@@ -78,7 +151,7 @@ def count_directories(dir_path, glob_pattern, recursive=False):
 
 
 def check_directory_exists(dir_path):
-    """Check if the directory exist."""
+    """Check if the directory exists."""
     if not os.path.exists(dir_path):
         raise ValueError(f"{dir_path} directory does not exist.")
     if not os.path.isdir(dir_path):
@@ -90,29 +163,24 @@ def create_directory(path: str, exist_ok=True) -> None:
     path = ensure_string_path(path, msg="'path' must be a string", accepth_pathlib=True)
     try:
         os.makedirs(path, exist_ok=exist_ok)
-        logger.debug(f"Created directory {path}.")
     except Exception as e:
+        dir_path = os.path.dirname(path)
         dir_name = os.path.basename(path)
-        msg = f"Can not create directory {dir_name} inside <path>. Error: {e}"
-        logger.exception(msg)
+        msg = f"Can not create directory {dir_name} inside {dir_path}. Error: {e}"
         raise FileNotFoundError(msg)
 
 
-def create_required_directory(dir_path, dir_name):
-    """Create directory <dir_name> inside the <dir_path> directory."""
-    try:
-        new_dir = os.path.join(dir_path, dir_name)
-        os.makedirs(new_dir, exist_ok=True)
-    except Exception as e:
-        msg = f"Can not create directory {dir_name} at {new_dir}. Error: {e}"
-        logger.exception(msg)
-        raise FileNotFoundError(msg)
+def create_required_directory(dir_path, dir_name, exist_ok=True):
+    """Create directory ``dir_name`` inside the ``dir_path`` directory."""
+    dir_path = ensure_string_path(dir_path, msg="'path' must be a string", accepth_pathlib=True)
+    new_dir_path = os.path.join(dir_path, dir_name)
+    create_directory(path=new_dir_path, exist_ok=exist_ok)
 
 
 def is_empty_directory(path):
     """Check if a directory path is empty.
 
-    Return False if path is a file or non-empty directory.
+    Return ``False`` if path is a file or non-empty directory.
     If the path does not exist, raise an error.
     """
     if not os.path.exists(path):
@@ -121,53 +189,47 @@ def is_empty_directory(path):
         return False
 
     paths = os.listdir(path)
-    if len(paths) == 0:
-        return True
-    else:
-        return False
+    return len(paths) == 0
 
 
-def _remove_file_or_directories(path):
-    """Return the file/directory or subdirectories tree of 'path'.
+def _remove_file_or_directories(path, logger=None):
+    """Return the file/directory or subdirectories tree of ``path``.
 
     Use this function with caution.
     """
     # If file
     if os.path.isfile(path):
         os.remove(path)
-        logger.info(f"Deleted the file {path}")
+        log_info(logger, msg=f"Deleted the file {path}")
     # If empty directory
     elif is_empty_directory(path):
         os.rmdir(path)
-        logger.info(f"Deleted the empty directory {path}")
+        log_info(logger, msg=f"Deleted the empty directory {path}")
     # If not empty directory
     else:
         shutil.rmtree(path)
-        logger.info(f"Deleted directories within {path}")
-    return None
+        log_info(logger, msg=f"Deleted directories within {path}")
 
 
-def remove_if_exists(path: str, force: bool = False) -> None:
-    """Remove file or directory if exists and force=True.
+def remove_if_exists(path: str, force: bool = False, logger=None) -> None:
+    """Remove file or directory if exists and ``force=True``.
 
-    If force=False --> Raise error
+    If ``force=False``, it raises an error.
     """
     # If the path does not exist, do nothing
     if not os.path.exists(path):
-        return None
+        return
 
     # If the path exists and force=False, raise Error
     if not force:
         msg = f"--force is False and a file already exists at: {path}"
-        logger.error(msg)
         raise ValueError(msg)
 
     # If force=True, remove the file/directory or subdirectories and files !
     try:
-        _remove_file_or_directories(path)
+        _remove_file_or_directories(path, logger=logger)
     except Exception as e:
         msg = f"Can not delete file(s) at {path}. The error is: {e}"
-        logger.error(msg)
         raise ValueError(msg)
 
 
@@ -181,18 +243,17 @@ def copy_file(src_filepath, dst_filepath):
         logger.info(msg)
     except Exception as e:
         msg = f"Something went wrong when copying {filename} into {dst_dir}.\n The error is: {e}."
-        logger.error(msg)
         raise ValueError(msg)
 
 
 def remove_path_trailing_slash(path: str) -> str:
-    """
+    r"""
     Removes a trailing slash or backslash from a file path if it exists.
 
     This function ensures that the provided file path is normalized by removing
-    any trailing directory separator characters ('/' or '\\'). This is useful for
-    maintaining consistency in path strings and for preparing paths for operations
-    that may not expect a trailing slash.
+    any trailing directory separator characters (``'/'`` or ``'\\'``).
+    This is useful for maintaining consistency in path strings and for
+    preparing paths for operations that may not expect a trailing slash.
 
     Parameters
     ----------

disdrodb/utils/encoding.py

@@ -0,0 +1,127 @@
+#!/usr/bin/env python3
+
+# -----------------------------------------------------------------------------.
+# Copyright (c) 2021-2023 DISDRODB developers
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+# -----------------------------------------------------------------------------.
+"""DISDRODB netCDF4 encoding utilities."""
+import xarray as xr
+
+EPOCH = "seconds since 1970-01-01 00:00:00"
+
+
+def set_encodings(ds: xr.Dataset, encoding_dict: dict) -> xr.Dataset:
+    """Apply the encodings to the xarray Dataset.
+
+    Parameters
+    ----------
+    ds  : xarray.Dataset
+        Input xarray dataset.
+    encoding_dict : dict
+        Dictionary with encoding specifications.
+
+    Returns
+    -------
+    xarray.Dataset
+        Output xarray dataset.
+    """
+    # Subset encoding dictionary
+    # - Here below encoding_dict contains only keys (variables) within the dataset
+    encoding_dict = {var: encoding_dict[var] for var in ds.data_vars if var in encoding_dict}
+
+    # Ensure chunksize smaller than the array shape
+    encoding_dict = sanitize_encodings_dict(encoding_dict, ds)
+
+    # Rechunk variables for fast writing !
+    # - This pop the chunksize argument from the encoding dict !
+    ds = rechunk_dataset(ds, encoding_dict)
+
+    # Set time encoding
+    ds["time"].encoding.update(get_time_encoding())
+
+    # Set the variable encodings
+    for var, encoding in encoding_dict.items():
+        ds[var].encoding.update(encoding)
+
+    # Ensure no deprecated "missing_value" attribute
+    # - When source dataset is netcdf (i.e. ARM)
+    for var in list(ds.variables):
+        _ = ds[var].encoding.pop("missing_value", None)
+
+    return ds
+
+
+def sanitize_encodings_dict(encoding_dict: dict, ds: xr.Dataset) -> dict:
+    """Ensure chunk size to be smaller than the array shape.
+
+    Parameters
+    ----------
+    encoding_dict : dict
+        Dictionary containing the variable encodings.
+    ds  : xarray.Dataset
+        Input dataset.
+
+    Returns
+    -------
+    dict
+        Encoding dictionary.
+    """
+    for var in ds.data_vars:
+        if var in encoding_dict:
+            shape = ds[var].shape
+            chunks = encoding_dict[var].get("chunksizes", None)
+            if chunks is not None:
+                chunks = [shape[i] if chunks[i] > shape[i] else chunks[i] for i in range(len(chunks))]
+                encoding_dict[var]["chunksizes"] = chunks
+    return encoding_dict
+
+
+def rechunk_dataset(ds: xr.Dataset, encoding_dict: dict) -> xr.Dataset:
+    """Coerce the dataset arrays to have the chunk size specified in the encoding dictionary.
+
+    Parameters
+    ----------
+    ds  : xarray.Dataset
+        Input xarray dataset
+    encoding_dict : dict
+        Dictionary containing the encoding to write the xarray dataset as a netCDF.
+
+    Returns
+    -------
+    xarray.Dataset
+        Output xarray dataset
+    """
+    for var in ds.data_vars:
+        if var in encoding_dict:
+            chunks = encoding_dict[var].pop("chunksizes", None)
+            if chunks is not None:
+                dims = list(ds[var].dims)
+                chunks_dict = dict(zip(dims, chunks))
+                ds[var] = ds[var].chunk(chunks_dict)
+    return ds
+
+
+def get_time_encoding() -> dict:
+    """Create time encoding.
+
+    Returns
+    -------
+    dict
+        Time encoding.
+    """
+    encoding = {}
+    encoding["units"] = EPOCH
+    encoding["calendar"] = "proleptic_gregorian"
+    return encoding

disdrodb/utils/list.py

@@ -0,0 +1,29 @@
+#!/usr/bin/env python3
+
+# -----------------------------------------------------------------------------.
+# Copyright (c) 2021-2023 DISDRODB developers
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+# -----------------------------------------------------------------------------.
+"""Utilities to work with lists."""
+
+
+def flatten_list(nested_list):
+    """Flatten a nested list into a single-level list."""
+    if isinstance(nested_list, list) and len(nested_list) == 0:
+        return nested_list
+    # If list is already flat, return as is to avoid flattening to chars
+    if isinstance(nested_list, list) and not isinstance(nested_list[0], list):
+        return nested_list
+    return [item for sublist in nested_list for item in sublist] if isinstance(nested_list, list) else [nested_list]