Convert a network graph from opareport (Intel Omni-Path) format to JSON

opa2json.py

#!/usr/bin/python
# -*- coding: utf-8 -*-

import sys
import json
import argparse
import xml.etree.ElementTree as et

import networkx as nx

parser = argparse.ArgumentParser(
    description="""Convert an XML network graph from opareport format to JSON.

This is meant as a way to share an Intel® Omni-Path Architecture (OPA)
network topology graph with the rest of the world, e.g.: D3.js or networkx.

The default behaviour reads the XML OPA graph from stdin
and writes the result to stdout:

$ opareport -o links -x | opa2json > opagraph.json

An output file can be specified:

$ opareport -o links -x | opa2json -o opagraph.json

An actual input file can be specified instead of stdin as well:

$ opareport -o links -x > opagraph.xml
$ opa2json opagraph.xml > opagraph.json

""",
    formatter_class=argparse.RawTextHelpFormatter,
)

parser.add_argument(
    "infile",
    nargs="?",
    type=argparse.FileType("r"),
    default=sys.stdin,
    help="""the input file containing the XML graph
(default: stdin)""",
)

parser.add_argument(
    "-o",
    "--out",
    nargs="?",
    metavar="PATH",
    type=argparse.FileType("w"),
    default=sys.stdout,
    help="""the output file where to write the JSON result to
(default: stdout)""",
)


def parse_element(element, *keys):
    attributes = {}
    for k in keys:
        v = element.find(k)
        if v is not None:
            attributes[k] = v.text
    return attributes


def parse(s_xml):
    parse_port = lambda e: parse_element(
        e, "NodeGUID", "PortNum", "NodeType", "NodeDesc"
    )
    parse_link = lambda e: parse_element(e, "Rate", "Rate_Int", "Internal")
    graph = nx.MultiGraph()
    root = et.fromstring(s_xml)
    for summary in root.findall("LinkSummary"):
        for link in summary.findall("Link"):
            # Nodes
            port_a, port_b = link.findall("Port")
            node_a = parse_port(port_a)
            node_b = parse_port(port_b)
            node_a_id = node_a.pop("NodeGUID")
            node_b_id = node_b.pop("NodeGUID")
            graph.add_node(node_a_id, **node_a)
            graph.add_node(node_b_id, **node_b)
            # Edge
            edge = parse_link(link)
            graph.add_edge(node_a_id, node_b_id, **edge)
    return graph


if __name__ == "__main__":
    args = parser.parse_args()
    s_xml = args.infile.read()
    graph = parse(s_xml)
    print(json.dumps(nx.readwrite.json_graph.node_link_data(graph)))
    # print(nx.drawing.nx_pydot.to_pydot(graph))