使用python查看文件夹占用情况

"""
计算指定目录大小
"""
import os
import sys
import prettytable as pt
import pandas as pd

LONGEST_FILE_PATH = 100


def classify_file_directory(path):
    """
    返回指定path下第一层文件和文件夹绝对路径列表
    param：
        path 文件夹路径
    returns：
        file_list_full: 第一层的文件绝对路径列表
        directory_list_full: 第一层文件夹的绝对路径列表
        file_list_short: 第一层文件名列表
        directory_list_short: 第一层文件绝对路径列表
    """
    # 第一层文件和目录
    mixed_file_directory = os.listdir(path)

    file_list_full = []
    directory_list_full = []

    file_list_short = []
    directory_list_short = []

    for mix in mixed_file_directory:
        full_path = os.path.join(path, mix)
        if os.path.isfile(full_path):
            file_list_full.append(full_path)
            file_list_short.append(mix)
        elif os.path.isdir(full_path):
            directory_list_full.append(full_path)
            directory_list_short.append(mix)

    return file_list_full, directory_list_full, file_list_short, directory_list_short


def get_directory_size(path):
    """
    获取文件夹大小
    """
    sum_size = 0
    try:
        for root, dirs, files in os.walk(path):
            for f in files:
                full_path = os.path.join(root, f)
                full_path = u"\\\\?\\" + full_path
                # 过滤掉路径过长的文件
                if len(full_path) > LONGEST_FILE_PATH:
                    continue
                size = os.path.getsize(full_path)
                sum_size += size
        return sum_size
    except Exception as err:
        print(err)


def get_file_size(path):
    """
    获取文件大小
    """
    return os.path.getsize(path)


def format_size(size_in_byte):
    """
    格式化大小，以B，KB，MB，GB结尾
    params:
        size_in_byte 以B为单位的文件空间大小
    returns：
        格式化后的以B、KB、MB、GB为单位的空间大小
    """
    precision = 1
    if size_in_byte < 1024:
        return str(round(size_in_byte, precision)) + "B"
    elif size_in_byte < 1024 ** 2:
        return str(round(size_in_byte / 1024, precision)) + "KB"
    elif size_in_byte < 1024 ** 3:
        return str(round(size_in_byte / 1024 ** 2, precision)) + "MB"
    else:
        return str(round(size_in_byte / 1024 ** 3, precision)) + "GB"


def get_pretty_table(path):
    """
    拼装prettyTable
    """
    file_list_full, directory_list_full, file_list_short, directory_list_short = classify_file_directory(path)

    file_size_list = [get_file_size(filePath) for filePath in file_list_full]
    directory_size_list = [get_directory_size(directory_path) for directory_path in directory_list_full]

    whole_name_list = []
    whole_size_list = []

    whole_name_list.extend(directory_list_short)
    whole_name_list.extend(file_list_short)

    whole_size_list.extend(directory_size_list)
    whole_size_list.extend(file_size_list)

    file_size_data = {"name": whole_name_list, "size": whole_size_list}
    file_size_df = pd.DataFrame(file_size_data)
    file_size_df.sort_values(by="size", inplace=True, ascending=False)
    
    tb = pt.PrettyTable()
    tb.add_column("name", list(file_size_df['name']), align='l')
    formated_file_size = [format_size(file_size_byte) for file_size_byte in list(file_size_df['size'])]
    tb.add_column("size", formated_file_size, align='l')
    return tb


if __name__ == "__main__":
    """
    python directory_detail.py D:\
    """
    pretty_table = get_pretty_table(sys.argv[1])
    print(pretty_table)