rhodecode/lib/rc_cache/archive_cache/utils.py

# Copyright (C) 2015-2024 RhodeCode GmbH
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License, version 3
# (only), as published by the Free Software Foundation.
#
# 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 Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#
# This program is dual-licensed. If you wish to learn more about the
# RhodeCode Enterprise Edition, including its added features, Support services,
# and proprietary license terms, please see https://rhodecode.com/licenses/

import os


class ArchiveCacheGenerationLock(Exception):
    pass


def archive_iterator(_reader, block_size: int = 4096 * 512):
    # 4096 * 64 = 64KB
    while 1:
        data = _reader.read(block_size)
        if not data:
            break
        yield data


def get_directory_statistics(start_path):
    """
    total_files, total_size, directory_stats = get_directory_statistics(start_path)

    print(f"Directory statistics for: {start_path}\n")
    print(f"Total files: {total_files}")
    print(f"Total size: {format_size(total_size)}\n")

    :param start_path:
    :return:
    """

    total_files = 0
    total_size = 0
    directory_stats = {}

    for dir_path, dir_names, file_names in os.walk(start_path):
        dir_size = 0
        file_count = len(file_names)

        for file in file_names:
            filepath = os.path.join(dir_path, file)
            file_size = os.path.getsize(filepath)
            dir_size += file_size

        directory_stats[dir_path] = {'file_count': file_count, 'size': dir_size}
        total_files += file_count
        total_size += dir_size

    return total_files, total_size, directory_stats


def format_size(size):
    # Convert size in bytes to a human-readable format (e.g., KB, MB, GB)
    for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
        if size < 1024:
            return f"{size:.2f} {unit}"
        size /= 1024