xgmml_networkx.py

__author__ = "Yasunobu OKAMURA, Michael P SCHROEDER"
__copyright__ = "Copyright (c) 2012 Y.Okamura"
__license__ = "GPL v3+"

import xml.parsers.expat
import networkx as nx

class XGMMLParserHelper(object):
    """
    Creates an instance of the XGMML parser that allows to import the
    network to a networkx.DiGraph object. Call instance.parseFile() and
    instance.graph() to import.
    """
    
    def __init__(self, graph=nx.DiGraph()):
        """
        
        Arguments:
        :param graph: Network X graph object
        :type graph: networkx.Graph
        """
        self._graph = graph
        self._parser = xml.parsers.expat.ParserCreate()
        self._parser.StartElementHandler = self._start_element
        self._parser.EndElementHandler = self._end_element
        self._tagstack = list()

        self._current_attr = dict()
        self._current_obj = dict()

    def _start_element(self, tag, attr):
        """
        
        Arguments:
        - `self`:
        - `tag`:
        - `attr`:
        """

        self._tagstack.append(tag)

        if tag == 'node' or tag == 'edge':
            self._current_obj = dict(attr)

        if tag == 'att' and (self._tagstack[-2] == 'node' or self._tagstack[-2] == 'edge'):
            if attr['type'] == 'string':
                if 'value' in attr:
                    self._current_attr[attr['name']] = attr['value']
            elif attr['type'] == 'real':
                if 'value' in attr:
                    self._current_attr[attr['name']] = float(attr['value'])
            elif attr['type'] == 'integer':
                if 'value' in attr:
                    self._current_attr[attr['name']] = int(attr['value'])
            elif attr['type'] == 'boolean':
                if 'value' in attr:
                    self._current_attr[attr['name']] = bool(attr['value'])
            else:
                raise NotImplementedError(attr['type'])

    def _end_element(self, tag):
        """
        
        Arguments:
        - `self`:
        - `tag`:
        """

        if tag == 'node':
            self._graph.add_node(self._current_obj['id'], label=self._current_obj['label'], **self._current_attr)
            #print 'add node', self._current_obj
        elif tag == 'edge':
            self._graph.add_edge(self._current_obj['source'], self._current_obj['target'], **self._current_attr)

        self._tagstack.pop()

    def parseFile(self, file):
        """
        :param file: The file the network will be written to. A file type object with the writing flag set.
        :type file: file
        """

        self._parser.ParseFile(file)

    def graph(self):
        """
        :rtype : networkx.DiGraph
        """

        return self._graph

def format_output(k,v):
    output_string = ""

    if type(v) in (str, unicode):
        return '<att name="{}" value="{}" type="string" />'.format(k, v)
    elif type(v) == float:
        return '<att name="{}" value="{}" type="real" />'.format(k, v)
    elif type(v) == int:
        return '<att name="{}" value="{}" type="integer" />'.format(k, v)
    elif type(v) == bool:
        if v:
            v = 1
        else:
            v = 0
        return '<att name="{}" value="{}" type="boolean" />'.format(k, v)
    else:
        raise NotImplementedError(str(type(v)) + ", " + str(k) + " " + str(v))

def XGMMLWriter(file, graph, graph_name):
    """


    :param file: The file the network will be written to. A file type object with the writing flag set.
    :type file: file
    :param graph: The graph that should be written to disk in the XGMML format.
    :type graph: networkx.DiGraph
    :param graph_name: Graph name.
    :type graph_name: str
    """

    print >>file, """<?xml version="1.0" encoding="UTF-8" standalone="yes"?>
<graph directed="1"  xmlns:dc="http://purl.org/dc/elements/1.1/" xmlns:xlink="http://www.w3.org/1999/xlink" xmlns:rdf="http://www.w3.org/1999/02/22-rdf-syntax-ns#" xmlns="http://www.cs.rpi.edu/XGMML">
<att name="selected" value="1" type="boolean" />
<att name="name" value="{0}" type="string"/>
<att name="shared name" value="{0}" type="string"/>""".format(graph_name)


    for onenode in graph.nodes(data=True):
        id = onenode[0]
        attr = dict(onenode[1])

        if 'label' in attr:
            label = attr['label']
            del attr['label']
        else:
            label = id
        
        print >>file, '<node id="{id}" label="{label}">'.format(id=id, label=label)
        for k, v in attr.iteritems():
            print >>file, format_output(k, v)
        print >>file, '</node>'
        
    for oneedge in graph.edges(data=True):
        print >>file, '<edge source="{}" target="{}">'.format(oneedge[0], oneedge[1])
        for k, v in oneedge[2].iteritems():
            print >>file, format_output(k, v)

        print >>file, '</edge>'
    print >>file, '</graph>'