view wsgraph/model.py @ 46:2d57cda87036

start noting useful graph DBs
author Jeff Hammel <jhammel@mozilla.com>
date Tue, 16 Apr 2013 01:03:19 -0700
parents d1e9602145fa
children
line wrap: on
line source


import sys
from abc import abstractmethod
from copy import deepcopy
from utils import isiterable

class Graph(object):
    """
    abstract base class for WSGraph model of a graph

    WSGraph is interacted with by implementing
    wsgraph.model.Graph object for a desired interface
    """

    @abstractmethod
    def node(self, name, value=None):
        """
        get or set a node

        When setting a node, a value of `None` will pop the value from
        the nodal values
        """
        # TODO: values should not be **kwargs as you could conceivably want
        # to set a node or edge to an empty dict
        # Instead, should be (self, name, values=None)

    @abstractmethod
    def nodes(self):
        """returns a list of all nodes"""

    @abstractmethod
    def edge(self, node1, node2, value=None):
        """
        get or set edge from node1 to node2
        """

    @abstractmethod
    def edges(self):
        """returns a list of all edges"""

    @abstractmethod
    def delete(self, node1, node2=None):
        """
        delete a node or edge if node2 is given
        If a node is deleted, all edges to it should be deleted
        """

    def __call__(self):
        """
        returns JSGN format of graph:
        {'nodes': {'node1': {},
                   'node2': {}, ...},
         'edges': {'node1': {'node2': {},
                             'node3': {}, ...},
                   'node2': {'node1': {}, ...}}
        }
        """
        retval = {'nodes': {}, 'edges': {}}
        for node in self.nodes():
            retval['nodes'][node] = self.node(node)
        for node1, node2 in self.edges():
            retval['edges'].setdefault(node1, {})[node2] = self.edge(node1, node2)
        return retval

    def __getitem__(self, key):
        """
        if key is a basestring, return the node of that name;
        if key is a 2-tuple/list, return the edge of that name
        """

        if isinstance(key, basestring) or (not isiterable(key)):
            return self.node(key)
        else:
            return self.edge(*key)

    def __setitem__(self, key, value):
        if isinstance(key, basestring) or (not isiterable(key)):
            self.node(key, value)
        else:
            key1, key2 = key
            self.edge(key1, key2, value)

    def __delitem__(self, key):
        if isinstance(key, basestring) or (not isiterable(key)):
            self.delete(key)
        else:
            key1, key2 = key
            self.delete(key1, key2)

    def __contains__(self, key):
        """
        if key is ..., returns if that node is in the graph
        if key is a 2-tuple/list, returns if the edge is in the graph
        """
        # XXX not necessarily the best implementation!
        if isinstance(key, basestring) or (not isiterable(key)):
            return key in self.nodes()
        else:
            return tuple(key) in self.edges()


class DirectedGraph(Graph):
    """mix-in class for directed graphs"""
    # TODO: is this possible without super or other black magicks?


class MemoryCache(Graph): # TODO -> MemoryCacheGraph
    """volatile in-memory representation of a graph"""

    # TODO: a subclass that pegs the type(s?) to something specific

    def __init__(self, node_type=dict):
        self._edges = {}
        self._nodes = {}
        self.node_type = node_type

    def node(self, name, value=None):
        if value is not None:
            # setter
            self._nodes[name] = deepcopy(value)
        else:
            # getter
            # TODO: deepcopy
            return deepcopy(self._nodes.get(name, None))

    def nodes(self):
        return self._nodes.keys()

    def edge(self, node1, node2, value=None):
        if value is not None:
            # setter
            self._edges[(node1, node2)] = deepcopy(value)
            for node in node1, node2:
                self._nodes.setdefault(node, self.node_type())
        else:
            # getter
            # TODO: deepcopy
            return deepcopy(self._edges.get((node1, node2), None))

    def edges(self):
        return self._edges.keys()

    def delete(self, node1, node2=None):
        if node2 is not None:
            # delete an edge
            key = node1, node2
            del self._edges[key]
            return

        # delete a node
        # TODO: if a node is deleted, all edges to it should be deleted
        del self._nodes[node1]


class FileCache(MemoryCache):
    """on-disk JSON file cache"""

    def __init__(self, filename):
        self.filename = filename
        raise NotImplementedError

# TODO: CLI entry point to convert from one model to another
# def main(args=sys.argv[1:]):