view wsgraph/model.py @ 8:f1f7a505e0d0

__contains__ method
author Jeff Hammel <jhammel@mozilla.com>
date Mon, 10 Dec 2012 17:08:02 -0800
parents cfcfa093e4b4
children 0affca1f4dc0
line wrap: on
line source

from abc import abstractmethod
from utils import iterable

class GraphModel(object):

    @abstractmethod
    def node(self, name, **values):
        """get or set a node"""

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

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

    @abstractmethod
    def edge(self, node1, node2, **values):
        """returns edge from node1 to node2"""

    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 iterable(key)):
            return self.node(key)
        else:
            return self.edge(*key)

    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 iterable(key)):
            return key in self.nodes()
        else:
            return tuple(key) in self.edges()

class MemoryCache(GraphModel):

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

    def node(self, name, **values):
        if values:
            # setter
            self._nodes[name] = values
        else:
            # getter
            return self._nodes.get(name, None)

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

    def edge(self, node1, node2, **values):
        if values:
            # setter
            self._edges[(node1, node2)] = values
        else:
            # getter
            return self._edges.get((node1, node2), None)

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