basant/basant/grab/tree.py

121 lines
3.0 KiB
Python

#! /usr/bin/env python
# -*- coding: utf-8 -*-
from icecream import ic
import json
"""network-vimwiki.py
Create site map for vimwiki.
Must run script under vimwiki directory.
Example:
output network map to network.json
dfsvw.py
Python version: 3"""
import sys
import re
linkRE = re.compile(r'\[\[([^\\\]|#]+)(?:#[^\\\]|]+)?(?:\|([^\\\]]+))?\]\]', re.UNICODE)
def probeLink(s):
"""
>>> probeLink(' abc')
[]
>>> probeLink(' [[]]')
[]
>>> probeLink(' [[abc]] [[abc|123]] ')
[('abc', ''), ('abc', '123')]
>>> a, b = probeLink(' [[中文1|中文2]]')[0]; print a, b
中文1 中文2
>>> probeLink(' [[abc|abc 123]]')
[('abc', 'abc 123')]
"""
ret = linkRE.findall(s)
return [(a, b) for a, b in ret]
def probeAllLink(fn):
"""extract all link in a wiki file"""
try:
lst = []
print(fn)
fp = open(fn)
for line in fp.readlines():
l = probeLink(line)
lst.extend(l)
except IOError:
print("cannot open",sys.stderr)
except UnicodeDecodeError:
lst = []
fp = open(fn, encoding="utf-16BE")
for line in fp.readlines():
l = probeLink(line)
lst.extend(l)
return lst
def outputNode(link, name, level):
"""output one wiki node, level for vimwiki tree hierarchy"""
""" Unused """
prefix = " " * 4 * level
if not name:
name = link
dct = {"prefix":prefix, "link":link, "name":name}
if sys.stdout.isatty():
fmt = "{prefix}{name}"
else:
if link == name:
fmt = "{prefix}[[{name}]]"
else:
fmt = "{prefix}[[{link}|{name}]]"
print(fmt.format(**dct))
def addNode(link, node, name, graph):
""" adds a node to the dictionary that tracks edges"""
if name:
node = name
if link in graph.keys():
graph[link].add(node)
else:
graph[link] = {node}
def bfs(index):
"""breath first search on vimwiki
keep seen and unseen to track which files to search
ugly loops instead of recursive setup """
seen = set()
unseen = [index]
graph = {}
filecount = 0
while 1:
try:
link = unseen.pop(0)
except IndexError:
print("Search complete...")
ic(filecount)
return graph
seen.add(link)
ret = probeAllLink(link + '.wiki')
for (node, name) in ret:
addNode(link, node, name, graph)
if not node in seen:
unseen.append(node)
filecount = filecount + 1
ic(filecount)
ic(seen)
def transformGraph(graph):
for key, value in graph.items():
graph[key] = list(value)
return graph
if __name__ == "__main__":
graph = bfs('index')
print("Transforming Graph...")
graph = transformGraph(graph)
print("Transform Complete...")
ic(graph)
print("Outputting file")
with open('network.json', 'w+') as out:
json.dump(graph, out)