矩形树图squarify

有一款磁盘分析神器space-sniffer，可以对磁盘空间进行可视化分析，我一直不知道这种图叫什么，搜了好久才弄清楚叫矩形树图（Treemap），而且查阅过程中还有个意外收获，即python中有一个专门的库squarify可以绘制这种图。

可用pip安装

pip install squarify

下面是个简单的示例

绘图代码如下

import squarify
import matplotlib.pyplot as plt

sizes  = [50, 30, 15, 5]
labels = ['A', 'B', 'C', 'D']
values = [500, 300, 150, 50]
colors = ['#FF9999', '#66B3FF', '#99FF99', '#FFCC99']

squarify.plot(
    sizes=sizes,
    label=labels,
    value=values,
    color=colors,
    alpha=0.8,
    pad=1,
    text_kwargs={'fontsize': 10, 'color': 'black', 'fontweight': 'bold'}
)

plt.axis('off')
plt.tight_layout()
plt.show()

该库的便捷之处在于，只需输入数据就可以出图，至于如何布局完全不用操心。其输入参数分别为

• 【sizes】图形尺寸
• 【labels】 标签
• 【values】 标签下面显示的数值
• 【colors】 颜色

分析磁盘数据

下面模仿space-sniffer的逻辑，做一个磁盘分析的小脚本，效果如下，这是一份电子书目录，通过矩形树图便捷地表现出了不同类别电子书的大小

代码如下，首先是递归查询文件夹的字节数，其逻辑很简单，如果子路径是文件，就累加其字节数，如果是文件夹，就调用自身，获取内部文件的累加字节数。

import os
from pathlib import Path

def getDirSize(path):
    total = 0
    for entry in os.scandir(path):
        if entry.is_file(follow_symlinks=False):
            total += entry.stat().st_size
        elif entry.is_dir(follow_symlinks=False):
            total += getDirSize(entry.path)
    return total

然后是绘制矩形树图的核心函数。

import squarify
import matplotlib.pyplot as plt
plt.rcParams['font.family'] = ['SimSun', 'Times New Roman']

def analyzeDir(root_path):
    root = Path(root_path).resolve()    
    sizes,labels = [], []
    for entry in os.scandir(root):
        size = getDirSize(entry.path)
        sizes.append(size)
        labels.append(f"{entry.name}\n({(size/1024**2):.1f} MB)")

    squarify.plot(
        sizes=sizes,
        label=labels,
        color=plt.cm.Set3(range(len(labels))),
        alpha=0.9,
        pad=0,
        text_kwargs={
            'fontsize': 9,
            'color': 'black',
            'fontweight': 'normal',
            'wrap': True
        }
    )
    plt.axis('off')
    plt.tight_layout()
    plt.show()    

if __name__ == "__main__":
    TARGET_PATH = r'E:\资源\01book\21信息技术基础'
    analyzeDir(TARGET_PATH)

其中，sizes即通过调用getDirSize获取的文件夹尺寸，labels中同样显示其字节数，以MB为单位。【plt.cm.Set3】可以为指定的序列分配颜色，在代码中可以为树状图中的每个矩形自动分配一组区分度高、视觉舒适的分类颜色。

小结

本文介绍了如何使用Python的squarify库绘制矩形树图(Treemap)，并实现磁盘空间分析功能。首先介绍了squarify库的安装和基本用法，通过简单的代码示例展示了如何绘制矩形树图。然后详细讲解了如何递归计算文件夹大小，并利用squarify库将磁盘目录结构可视化，最终生成类似space-sniffer工具的磁盘分析图表。文章提供了完整的代码实现，包括递归计算文件夹大小的函数和绘制矩形树图的核心函数，帮助读者快速实现磁盘空间可视化分析功能。