sl-shared-assets 6.1.1__py3-none-any.whl

sl_shared_assets/data_transfer/checksum_tools.py

@@ -0,0 +1,125 @@
+"""This module provides the assets for computing data integrity checksums."""
+
+import os
+from pathlib import Path
+from functools import partial
+from concurrent.futures import ProcessPoolExecutor, as_completed
+
+from tqdm import tqdm
+import xxhash
+
+# Defines a 'blacklist' set of files. Primarily, this list contains the service files that may change after the session
+# data has been acquired. Therefore, it does not make sense to include them in the checksum, as they do not reflect the
+# data that should remain permanently unchanged.
+_excluded_files = {
+    "ax_checksum.txt",
+    "nk.bin",
+}
+
+
+def _calculate_file_checksum(base_directory: Path, file_path: Path) -> tuple[str, bytes]:  # pragma: no cover
+    """Calculates the xxHash3-128 checksum for the target file and its path relative to the base directory.
+
+    Args:
+        base_directory: The path to the directory being processed by the 'calculate_directory_checksum' function.
+        file_path: The path to the target file relative to the base directory.
+
+    Returns:
+        A tuple with two elements. The first element is the path to the file relative to the base directory. The second
+        element is the xxHash3-128 checksum that reflects the file's path and data.
+    """
+    # Initializes the hashsum object.
+    checksum = xxhash.xxh3_128()
+
+    # Encodes the relative path and appends it to the checksum. This ensures that the hashsum reflects both the state
+    # of individual files and the layout of the overall encoded directory structure.
+    relative_path = str(file_path.relative_to(base_directory))
+    checksum.update(relative_path.encode())
+
+    # Extends the checksum to reflect the file data state. Uses 8 MB chunks to avoid excessive RAM hogging at the cost
+    # of slightly reduced throughput.
+    with file_path.open("rb") as f:
+        for chunk in iter(lambda: f.read(1024 * 1024 * 8), b""):
+            checksum.update(chunk)
+
+    # Returns both path and file checksum. Although the relative path information is already encoded in the hashsum, the
+    # relative path information is re-encoded at the directory level to protect against future changes to the per-file
+    # hashsum calculation logic. It is extra work, but it improves the overall checksum security.
+    return relative_path, checksum.digest()
+
+
+def calculate_directory_checksum(
+    directory: Path, num_processes: int | None = None, *, progress: bool = False, save_checksum: bool = True
+) -> str:
+    """Calculates the xxHash3-128 checksum for the input directory.
+
+    Note:
+        The function can be configured to write the generated checksum as a hexadecimal string to the ax_checksum.txt
+        file stored at the highest level of the input directory.
+
+        The xxHash3 checksum is not suitable for security purposes and is only used to ensure data integrity.
+
+        The returned checksum accounts for both the contents of each file and the layout of the input directory
+        structure.
+
+    Args:
+        directory: The path to the directory for which to generate the checksum.
+        num_processes: The number of processes to use for parallelizing checksum calculation. If set to None, the
+            function uses all available CPU cores.
+        progress: Determines whether to track the checksum calculation progress using a progress bar.
+        save_checksum: Determines whether the checksum should be saved (written to) a .txt file.
+
+    Returns:
+        The xxHash3-128 checksum for the input directory as a hexadecimal string.
+    """
+    # Determines the number of parallel processes to use.
+    if num_processes is None:
+        num_processes = max(1, os.cpu_count())  # type: ignore[type-var]
+
+    # Determines the path to each file inside the input directory structure and sorts them for consistency
+    files = sorted(
+        path
+        for path in directory.rglob("*")
+        if path.is_file() and f"{path.stem}{path.suffix}" not in _excluded_files  # Excludes service files
+    )
+
+    # Pre-creates the directory checksum
+    checksum = xxhash.xxh3_128()
+
+    # Process files in parallel
+    with ProcessPoolExecutor(max_workers=num_processes) as executor:
+        # Creates the partial function with fixed base_directory (the first argument of _calculate_file_hash())
+        process_file = partial(_calculate_file_checksum, directory)
+
+        # Submits all tasks to be executed in parallel
+        # noinspection PyTypeChecker
+        future_to_path = {executor.submit(process_file, file): file for file in files}
+
+        # Collects results as they complete
+        results = []
+        if progress:
+            with tqdm(
+                total=len(files), desc=f"Calculating checksum for {Path(*directory.parts[-4:-1])}", unit="file"
+            ) as pbar:
+                for future in as_completed(future_to_path):
+                    results.append(future.result())
+                    pbar.update(1)
+        else:
+            # For batch mode, uses a direct list comprehension with as_completed. This avoids the overhead of progress
+            # tracking while maintaining parallel processing, avoiding terminal clutter in batched contexts.
+            results = [future.result() for future in as_completed(future_to_path)]
+
+        # Sorts results for consistency and combines them into the final checksum
+        for file_path, file_checksum in sorted(results):
+            checksum.update(file_path.encode())
+            checksum.update(file_checksum)
+
+    checksum_hexstring = checksum.hexdigest()
+
+    # Writes the hash to ax_checksum.txt in the root directory
+    if save_checksum:
+        checksum_path = directory.joinpath("ax_checksum.txt")
+        with checksum_path.open("w") as f:
+            f.write(checksum_hexstring)
+
+    return checksum_hexstring

sl_shared_assets/data_transfer/transfer_tools.py

@@ -0,0 +1,181 @@
+"""This module provides the assets for moving the data between destinations available on the host-machine's filesystem
+and removing the data from the host-machine.
+"""
+
+import os
+import shutil
+from pathlib import Path
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
+from tqdm import tqdm
+from ataraxis_time import PrecisionTimer
+from ataraxis_base_utilities import console, ensure_directory_exists
+
+from .checksum_tools import calculate_directory_checksum
+
+
+def delete_directory(directory_path: Path) -> None:
+    """Deletes the target directory and all its subdirectories using parallel processing.
+
+    Args:
+        directory_path: The path to the directory to delete.
+    """
+    # Checks if the directory exists and, if not, aborts early
+    if not directory_path.exists():
+        return
+
+    # Builds the list of files and directories inside the input directory using Path
+    files = [p for p in directory_path.iterdir() if p.is_file()]
+    subdirectories = [p for p in directory_path.iterdir() if p.is_dir()]
+
+    # Deletes files in parallel
+    with ThreadPoolExecutor() as executor:
+        list(executor.map(os.unlink, files))  # Forces completion of all tasks
+
+    # Recursively deletes subdirectories
+    for subdir in subdirectories:
+        delete_directory(subdir)
+
+    # Removes the now-empty root directory. Since Windows is sometimes slow to release file handles, adds
+    # an optional delay step to give Windows time to release file handles.
+    max_attempts = 5
+    delay_timer = PrecisionTimer("ms")
+    for _ in range(max_attempts):
+        # noinspection PyBroadException
+        try:
+            directory_path.rmdir()
+            break  # Breaks early if the call succeeds
+        except Exception:  # pragma: no cover
+            delay_timer.delay(block=False, delay=500, allow_sleep=True)  # For each failed attempt, sleeps for 500 ms
+            continue
+
+
+def _transfer_file(source_file: Path, source_directory: Path, destination_directory: Path) -> None:
+    """Copies the input file from the source directory to the destination directory while preserving the file metadata.
+
+    This worker method is used by the transfer_directory() method to move multiple files in parallel.
+
+    Notes:
+        If the file is found under a hierarchy of subdirectories inside the input source_directory, that hierarchy will
+        be preserved in the destination directory.
+
+    Args:
+        source_file: The file to be copied.
+        source_directory: The root directory where the file is located.
+        destination_directory: The destination directory where to move the file.
+    """
+    relative = source_file.relative_to(source_directory)
+    dest_file = destination_directory / relative
+    shutil.copy2(source_file, dest_file)
+
+
+def transfer_directory(
+    source: Path,
+    destination: Path,
+    num_threads: int = 1,
+    *,
+    verify_integrity: bool = False,
+    remove_source: bool = False,
+    progress: bool = False,
+) -> None:
+    """Copies the contents of the input source directory to the destination directory while preserving the underlying
+    directory hierarchy.
+
+    Notes:
+        This function recreates the moved directory hierarchy on the destination if the hierarchy does not exist. This
+        is done before copying the files.
+
+        The function executes a multithreaded copy operation and does not by default remove the source data after the
+        copy is complete.
+
+        If the function is configured to verify the transferred data's integrity, it generates an xxHash-128 checksum of
+        the data before and after the transfer and compares the two checksums to detect data corruption.
+
+    Args:
+        source: The path to the directory to be transferred.
+        destination: The path to the destination directory where to move the contents of the source directory.
+        num_threads: The number of threads to use for the parallel file transfer. Setting this value to a number below
+            1 instructs the function to use all available CPU threads.
+        verify_integrity: Determines whether to perform integrity verification for the transferred files.
+        remove_source: Determines whether to remove the source directory after the transfer is complete and
+            (optionally) verified.
+        progress: Determines whether to track the transfer progress using a progress bar.
+
+    Raises:
+        RuntimeError: If the transferred files do not pass the xxHas3-128 checksum integrity verification.
+    """
+    if not source.exists():
+        message = f"Unable to transfer the source directory {source}, as it does not exist."
+        console.error(message=message, error=FileNotFoundError)
+
+    # If the number of threads is less than 1, interprets this as a directive to use all available CPU cores.
+    if num_threads < 1:
+        cpu_count = os.cpu_count()
+        num_threads = cpu_count if cpu_count is not None else 1
+
+    # If transfer integrity verification is enabled, but the source directory does not contain the 'ax_checksum.txt'
+    # file, checksums the directory before the transfer operation.
+    if verify_integrity and not source.joinpath("ax_checksum.txt").exists():
+        calculate_directory_checksum(directory=source, progress=False, save_checksum=True)
+
+    # Ensures the destination root directory exists.
+    ensure_directory_exists(destination)
+
+    # Collects all items (files and directories) in the source directory.
+    all_items = tuple(source.rglob("*"))
+
+    # Loops over all items (files and directories). Adds files to the file_list variable. Uses directories to reinstate
+    # the source subdirectory hierarchy in the destination directory.
+    file_list = []
+    for item in sorted(all_items, key=lambda x: len(x.relative_to(source).parts)):
+        # Recreates directory structure on destination
+        if item.is_dir():
+            dest_dir = destination / item.relative_to(source)
+            dest_dir.mkdir(parents=True, exist_ok=True)
+        # Also builds the list of files to be moved
+        else:  # is_file()
+            file_list.append(item)
+
+    # Copies the data to the destination. For parallel workflows, the method uses the ThreadPoolExecutor to move
+    # multiple files at the same time. Since I/O operations do not hold GIL, we do not need to parallelize with
+    # Processes here.
+    if num_threads > 1:
+        with ThreadPoolExecutor(max_workers=num_threads) as executor:
+            futures = {executor.submit(_transfer_file, file, source, destination): file for file in file_list}
+            for future in tqdm(
+                as_completed(futures),
+                total=len(file_list),
+                desc=f"Transferring files to {Path(*destination.parts[-6:])}",
+                unit="file",
+                disable=not progress,
+            ):
+                # Propagates any exceptions from the file transfer.
+                future.result()
+    else:
+        for file in tqdm(
+            file_list,
+            desc=f"Transferring files to {Path(*destination.parts[-6:])}",
+            unit="file",
+            disable=not progress,
+        ):
+            _transfer_file(file, source, destination)
+
+    # Verifies the integrity of the transferred directory by rerunning xxHash3-128 calculation.
+    if verify_integrity:
+        destination_checksum = calculate_directory_checksum(directory=destination, progress=False, save_checksum=False)
+        with source.joinpath("ax_checksum.txt").open("r") as local_checksum:
+            if destination_checksum != local_checksum.readline().strip():
+                message = (
+                    f"Checksum mismatch detected when transferring {Path(*source.parts[-6:])} to "
+                    f"{Path(*destination.parts[-6:])}! The data was likely corrupted in transmission."
+                )
+                console.error(message=message, error=RuntimeError)
+
+    # If necessary, removes the transferred directory from the original location.
+    if remove_source:
+        message = (
+            f"Removing the now-redundant source directory {source} and all of its contents following the successful "
+            f"transfer..."
+        )
+        console.echo(message=message)
+        delete_directory(source)