python项目结构图_用python+graphviz/networkx画目录结构树状图

想着用python绘制某目录的树状图，一开始想到了用grapgviz，因为去年离职的时候整理文档，用graphviz画过代码调用结构图。graphviz有一门自己的语言DOT，dot很简单，加点加边设置属性就这点东西，而且有python接口。

我在ubuntu下，先要安装graphviz软件，官网有deb包，然后python安装pygraphviz模块。

目标功能是输入一个路径，输出该路径下的目录结构以及文件和文件夹的大小

# -*- coding:utf-8 -*-

import pygraphviz as pgv

import os

from os.path import getsize, join

import sys

import Queue

reload(sys)

sys.setdefaultencoding('utf8')

class DrawDirectorTree():

"""

绘制目录结构图，用树的形式

"""

@classmethod

def getdirsize(cls, dir):

"""

获取文件夹大小

:param dir:

:return: 返回尺寸

"""

size = 0L

for root, dirs, files in os.walk(dir):

size += sum([getsize(join(root, name)) for name in files])

return size

@classmethod

def draw_director_tree(cls, input_path):

"""

深度遍历一个目录，绘制目录树形图

:param input_path: 目标目录

:return:

"""

if (not os.path.exists(input_path)) or (not os.path.isdir(input_path)):

print "Input_path Error!"

return None

# 用队列做BFS

director_queue = Queue.Queue()

director_queue.put(input_path)

# 初始化一个图

tree_graph = pgv.AGraph(directed=True, strict=True)

tree_graph.node_attr['style'] = 'filled'

tree_graph.node_attr['shape'] = 'square'

tree_graph.add_node(input_path + "\n" + str(os.path.getsize(input_path)))

while not director_queue.empty():

new_parent = director_queue.get()

if os.path.isdir(new_parent):

child_list = os.listdir(new_parent)

for child in child_list:

full_child = join(new_parent, child)

if os.path.isfile(full_child):

new_parent_lable = new_parent + "\n" + str(os.path.getsize(new_parent))

child_lable = full_child + "\n" + str(os.path.getsize(full_child))

tree_graph.add_node(child_lable)

tree_graph.add_edge(new_parent_lable, child_lable)

elif os.path.isdir(full_child):

new_parent_lable = new_parent + "\n" + str(os.path.getsize(new_parent))

child_lable = full_child + "\n" + str(os.path.getsize(full_child))

tree_graph.add_node(child_lable)

tree_graph.add_edge(new_parent_lable, child_lable)

director_queue.put(full_child)

tree_graph.graph_attr['epsilon'] = '0.001'

print tree_graph.string() # print dot file to standard output

tree_graph.write('director_tree.dot')

tree_graph.layout('dot') # layout with dot

tree_graph.draw('director_tree.png') # write to file

DrawDirectorTree.draw_director_tree('/home/aron/workspace/python_space/')

运行会生成一个dot文件，就是有向图的结构文件，还有存为一个png图

事实上效果不太好，小的目录可以搞定，大的目录就耗时太长了，主要是计算文件夹尺寸本身就很耗时，而且无法生成巨大的结构图。

之后发现networkx模块的使用方式几乎与graphviz一样，add_node、add_edge，就写了一版networkx的。

# -*- coding:utf-8 -*-

import os

from os.path import getsize, join

import sys

import Queue

import networkx as nx

import matplotlib.pyplot as plt

reload(sys)

sys.setdefaultencoding('utf8')

class DirectorTreeNetworkx():

"""

用networkx绘制目录结构图

"""

@classmethod

def getdirsize(cls, dir):

"""

获取文件夹大小

:param dir:

:return: 返回尺寸

"""

size = 0L

for root, dirs, files in os.walk(dir):

size += sum([getsize(join(root, name)) for name in files])

return size

@classmethod

def draw_director_tree(cls, input_path):

"""

深度遍历一个目录，绘制目录树形图

:param input_path: 目标目录

:return:

"""

if (not os.path.exists(input_path)) or (not os.path.isdir(input_path)):

print "Input_path Error!"

return None

# 用队列做BFS

director_queue = Queue.Queue()

director_queue.put(input_path)

# 初始化一个图

tree_graph = nx.DiGraph()

tree_graph.add_node(input_path + "\n" + str(os.path.getsize(input_path)))

while not director_queue.empty():

new_parent = director_queue.get()

if os.path.isdir(new_parent):

child_list = os.listdir(new_parent)

for child in child_list:

full_child = join(new_parent, child)

if os.path.isfile(full_child):

new_parent_lable = new_parent + "\n" + str(os.path.getsize(new_parent))

child_lable = full_child + "\n" + str(os.path.getsize(full_child))

tree_graph.add_node(child_lable)

tree_graph.add_edge(new_parent_lable, child_lable)

elif os.path.isdir(full_child):

new_parent_lable = new_parent + "\n" + str(os.path.getsize(new_parent))

child_lable = full_child + "\n" + str(os.path.getsize(full_child))

tree_graph.add_node(child_lable)

tree_graph.add_edge(new_parent_lable, child_lable)

director_queue.put(full_child)

nx.draw(tree_graph)

plt.show()能够解决稍大一点的目录，主要是生成网络图片比graphviz快一点，能展示稍微复杂一点的网络结构比如

这个也就玩一玩，真正做目录结构展示还得想别的办法，同时不能期望在获取结构的同时计算文件夹大小。

上面的代码计算文件夹大小的策略也有点问题，存在重复计算，算一次就够了。同时还可以先构建完结构自底向上的计算目录大小，这样就避免每计算一层目录都要walk一遍求和了，下层文件大小只计算一次就行了。