RFCartography/rfcartography/rfcartographer.py
2023-01-03 14:42:54 +01:00

173 lines
8.0 KiB
Python

from networkx import MultiDiGraph, kamada_kawai_layout, \
draw_networkx_nodes, draw_networkx_edges
from matplotlib.pyplot import figure, close
from matplotlib.figure import Figure
from matplotlib.layout_engine import TightLayoutEngine
from io import StringIO
from math import sqrt
from rfcartography.index_parser import Document, DocType
class RFCMap:
def __init__(self,
graph: MultiDiGraph,
nodes: dict[DocType, list[Document]],
edges: dict[str, tuple[Document, Document]],
url_base: str,
node_color: dict[DocType, str] = {DocType.RFC: '#2072b1',
DocType.STD: '#c21a7e',
DocType.BCP: '#6d388d',
DocType.FYI: '#8bbd3e',
DocType.NIC: '#efe50b',
DocType.IEN: '#f28f20',
DocType.RTR: '#e32326'},
edge_style: dict[str, tuple[str, str]]= {'obsoletes': ('dashed', '#607d8d'),
'obsoleted by': ('dashed', '#303c50'),
'updates': ('dashdot', '#607d8d'),
'updated by': ('dashdot', '#303c50'),
'is also': ('solid', '#132e41'),
'see also': ('dotted', '#008e90')}):
self.graph: MultiDiGraph = graph
self.nodes: dict[DocType, list[Document]] = nodes
self.edges: dict[str, tuple[Document, Document]] = edges
self.node_color: dict[DocType, str] = node_color
self.edge_style: dict[str, tuple[str, str]] = edge_style
self.url_base: str = url_base
self.position: dict = kamada_kawai_layout(self.graph)
return
def set_node_color(self,
doctype: DocType,
color: str) -> None:
self.node_color[doctype] = color
return
def set_edge_style(self,
reftype: str,
style: tuple[str, str]) -> None:
self.edge_style[reftype] = style
return
def set_url_base(self,
url_base: str) -> None:
self.url_base = url_base
return
def get_node_colors(self) -> dict[DocType, str]:
return self.node_color
def get_edge_styles(self) -> dict[str, tuple[str, str]]:
return self.edge_style
def get_url_base(self) -> str:
return self.url_base
def get_node_count(self) -> int:
return self.graph.number_of_nodes()
def get_edge_count(self) -> int:
return self.graph.size()
def draw(self) -> str:
nodes: list[Document] = []
node_colors: list[str] = []
urls: list[str] = []
for doctype in self.nodes:
nodes = nodes + self.nodes[doctype]
node_colors = node_colors + [self.node_color[doctype]]*len(self.nodes[doctype])
for node in self.nodes[doctype]:
urls.append(f"{self.url_base}{node.docID()}")
edge_rad: dict[float] = {'obsoletes': 0.1 if len(self.edges['obsoleted by']) > 0 else 0,
'obsoleted by': 0.1 if len(self.edges['obsoletes']) > 0 else 0,
'updates': 0.1 if len(self.edges['updated by']) > 0 else 0,
'updated by': 0.1 if len(self.edges['updates']) > 0 else 0,
'is also': 0,
'see also': 0}
size: float = sqrt(len(nodes))
fig: Figure = figure(figsize=(size, size), layout=TightLayoutEngine(pad=0.2))
fig.clear()
draw_networkx_nodes(self.graph, self.position,
nodelist=nodes,
node_size=128,
node_color=node_colors,
node_shape="o").set_urls(urls)
for reftype in self.edges:
draw_networkx_edges(self.graph, self.position,
edgelist=self.edges[reftype],
connectionstyle=f"arc3,rad={edge_rad[reftype]}",
style=self.edge_style[reftype][0],
edge_color=self.edge_style[reftype][1])
svg: StringIO = StringIO()
fig.savefig(svg, format='svg')
svg.seek(0)
close(fig)
return svg.read().removeprefix('<?xml version="1.0" encoding="utf-8" standalone="no"?>\n<!DOCTYPE svg PUBLIC "-//W3C//DTD SVG 1.1//EN"\n "http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd">\n')
class RFCartographer:
def __init__(self,
index: dict[DocType: dict[int, Document]]):
self.index: dict[DocType: dict[int, Document]] = index
return
def get_document(self,
doctype: DocType,
number: int) -> Document | None:
return self.index[doctype].get(number, None)
def map_subnet(self,
core: Document,
url: str,
max_depth: int = 0,
node_color: dict[DocType, str] = None,
edge_style: dict[str, tuple[str, str]] = None,
node_types: list[DocType] = []) -> RFCMap:
"""generate a map for the subnet core belongs to"""
if node_types == []:
node_types = [DocType.RFC, DocType.STD, DocType.BCP,
DocType.FYI, DocType.NIC, DocType.IEN, DocType.RTR]
nodes: dict = {DocType.RFC: [],
DocType.STD: [],
DocType.BCP: [],
DocType.FYI: [],
DocType.NIC: [],
DocType.IEN: [],
DocType.RTR: []}
edges: dict = {'obsoletes': [],
'obsoleted by': [],
'updates': [],
'updated by': [],
'is also': [],
'see also': []}
params: dict[str, dict] = {}
if node_color is not None:
params['node_color'] = node_color
if edge_style is not None:
params['edge_style'] = edge_style
todo: list[tuple[Document, int]] = [(core, 0)]
done: list[Document] = []
graph: MultiDiGraph = MultiDiGraph()
graph.add_node(core)
nodes[core.type].append(core)
while len(todo) > 0:
node: tuple[Document, int] = todo.pop(0)
if node[0] not in done:
done.append(node[0])
if node[1] < max_depth or max_depth <= 0:
for neighbor in node[0].get_references():
if not neighbor[1].type in node_types:
continue
if not graph.has_node(neighbor[1]):
graph.add_node(neighbor[1])
nodes[neighbor[1].type].append(neighbor[1])
graph.add_edge(node[0], neighbor[1], reftype=neighbor[0])
edges[neighbor[0]].append((node[0], neighbor[1]))
todo.append((neighbor[1], node[1]+1))
return RFCMap(graph, nodes, edges, url, **params)