Как стать автором
Обновить

Визуализация каталогов на Python средствами NetworkX

Python *
Из песочницы
Листая на Хабре раздел Python наткнулся на интересную статью о библиотеке NetworkX. Впечатлившись красивыми графами, решил повысить свой python-скилл и покопаться в networkx.
image

Пролог


Первый вопрос — откуда взять данные для визуализации? Генерировать случайные не интересно, они и в комплекте модуля были. Тут вспомнилась Dos утилитка tree, выводящая каталоги файловой системы в виде дерева. Решено было написать красивый аналог на Python и нарисовать все в networkx с помощью matplotlib.

Акт первый


Copy Source | Copy HTML
  1. def get_tree(tree=[u"E:\\Музыка",], G=nx.Graph(), itr= 0, max_itr=900):
  2.     point = tree.pop( 0)
  3.     itr = itr + 1
  4.     sub_tree = [os.path.join(point, x) for x in os.listdir(point) if os.path.isdir(os.path.join(point, x)) and not is_hidden_dir(os.path.join(point, x))]
  5.     if sub_tree:
  6.         tree.extend(sub_tree)
  7.         G.add_edges_from(map(lambda b : (point, b), sub_tree))
  8.     if tree and itr <= max_itr:
  9.         return get_tree(tree, G, itr)
  10.     else:
  11.         return G

Тут уже становится понятно, что музыки не так уж и много, судя по получившемуся графу.
Сам по себе код не представляет ничего сложно. Функция рекурсивно вызывает сама себя пока не закончится список директорий на разбор tree или пока не выйдет установленное количество итераций. При каждом вызове из списка директорий «выталкивается» первый элемент, у которого ищутся подкаталоги стандартной python функцией os.listdir(). После чего списки сливаются и полученные связи добавляются в структуру графа. Главной хитростью является функция is_hidden_dir(). Она проверяет, является ли файл скрытым. Вначале я думал, что эту тривиальную задачу можно решить средствами самого языка. Но оказалось, что это прост только в Unix системах, а в Windows, под которой я пишу, эта задача приобретает оттенки мазохизма.

Акт второй


Copy Source | Copy HTML
  1. def is_hidden_dir(d):
  2.     import sys, subprocess
  3.     if sys.platform.startswith("win"):
  4.         p = subprocess.check_output(["attrib", d])
  5.         return True if 'H' in p[:12] else False
  6.     else:
  7.         return True if os.path.basename(d)[ 0] == '.' else False


Сначала проверяем систему, пользователя. Если не Windows, то определение свойства hidden у файла банально — первым символом каталога должна быть точка «.».
В случае винды все нетипичнее. Подумав как можно сделать нечто подобное, я понял, что явных свойств у файла как точка в nix'ах не будет. Разбираться с WinAPI под Python не было желания. Оставалось воспользоваться грязными хаками, что я и сделал.

Быстрым поиском в гугле удалось найти консольную утилитку attrib. Изначально планировалось просто запустить ее через os.system(), но документация меня отговорила. Тем более мне был важен вывод утилиты, а не сам факт ее успешной работы. В манах нашлась нужная функция subprocess.check_output(), возвращающая результат работы аргумента. Дальнейшая проверка проста, в первых 12 символах надо найти вхождение флага «H». Но и тут виндовс не дал мне расслабиться. Пока он обрабатывал директории, в именах которых присутствовали только латинские символы, все было спокойно, но когда он добрался до русских букв, то работать отказался наотрез. Логично, что его не устраивала кодировка. Я с легким сердцем вставил декодер
Copy Source | Copy HTML
  1. p = subprocess.check_output(["attrib", d.encode('cp1251')])

Но питон быстро переубедил меня в успешности такой идее, отказавшись на этот раз работать с англоязычными названиями.
Итогом стало компромиссное решение.
Copy Source | Copy HTML
  1. p = subprocess.check_output(["attrib", d.encode('cp1251') if isinstance(d, unicode) else d])


Акт третий. Заключительный


Осталось лишь нарисовать сгенерированный граф на networkx, ради чего все и затевалось. Не надо долго гадать, что бы понять, что и тут надо было поковыряться. Функция отрисовки заманчиво проста:
Copy Source | Copy HTML
  1. import networkx as nx
  2. import matplotlib.pyplot as plt
  3.  
  4. def main():
  5.     G = get_tree()
  6.     nx.draw(G, with_labels=False, node_color="blue", alpha= 0.6, node_size=50)
  7.     plt.savefig("edge_colormap.png")
  8.     plt.show()

Параметры указательные в nx.draw() не обязательны, кроме самого графа. NetwokX поддерживает 2 библиотеки отрисовки визуальных данных: matplotlib и pygraphviz. Сначала я решил использовать pygraphviz. Скачал его с оффсайта, установил, начал устанавливать для него обертку под питон pygraphviz, но тут pip ругнулся и сказал, что pygraphviz отказывается иметь дело с windows. Ладно, подумал я, у нас есть альтернатива. Matplotlib встал без лишних вопросов, но при запуске скрипта с графом начал возмущаться, что я до сих пор не пользуюсь NumPy. Скачали и поставил NumPy. Он ничего у меня не попросил и просто начал работать.

Итоговый код выглядит так:
Copy Source | Copy HTML
  1. #-*- encoding: utf-8
  2.  
  3. import networkx as nx
  4. import matplotlib.pyplot as plt
  5. import os
  6.  
  7. def get_tree(tree=[u"E:\\Музыка",], G=nx.Graph(), itr= 0, max_itr=900):
  8.     point = tree.pop( 0)
  9.     itr = itr + 1
  10.     sub_tree = [os.path.join(point, x) for x in os.listdir(point) if os.path.isdir(os.path.join(point, x)) and not is_hidden_dir(os.path.join(point, x))]
  11.     if sub_tree:
  12.         tree.extend(sub_tree)
  13.         G.add_edges_from(map(lambda b : (point, b), sub_tree))
  14.     if tree and itr <= max_itr:
  15.         return get_tree(tree, G, itr)
  16.     else:
  17.         return G
  18.  
  19. def is_hidden_dir(d):
  20.     import sys, subprocess
  21.     if sys.platform.startswith("win"):
  22.         p = subprocess.check_output(["attrib", d.encode('cp1251') if isinstance(d, unicode) else d])
  23.         return True if 'H' in p[:12] else False
  24.     else:
  25.         return True if os.path.basename(d)[ 0] == '.' else False
  26.  
  27. def main():
  28.     G = get_tree()
  29.     nx.draw(G, with_labels=False, node_color="blue", alpha= 0.6, node_size=50)
  30.     plt.savefig("edge_colormap.png")
  31.     plt.show()
  32.  
  33. if __name__ == "__main__":
  34.     main()


Занавес


Снимок всей системы получить так и не удалось. Памяти не хватает или рук, что вероятнее.
Несколько картинок напоследок:
<img title="" border=«0» alt="" src="
Теги:
Хабы:
Всего голосов 49: ↑48 и ↓1 +47
Просмотры 14K
Комментарии Комментарии 35

Работа

Data Scientist
115 вакансий
Python разработчик
199 вакансий