Subversion Repositories programming

#!/usr/bin/env python

__author__    = "Ira W. Snyder (devel@irasnyder.com)"
__copyright__ = "Copyright (c) 2006 Ira W. Snyder (devel@irasnyder.com)"
__license__   = "GNU GPL v2 (or, at your option, any later version)"

from PyCompat import * # fixes for school computers

class Vertex (object):
        def __init__ (self, label='', value=0, shape='circle', raw_obj=None):
                self.label = label
                self.value = value
                self.shape = shape
                self.raw_obj = raw_obj

class Edge (object):
        def __init__ (self, color='black', label=''):
                self.color = color
                self.label = label

class Graph:
        """Class that will handle weighted, undirected graphs. It can have
           arbitrary data stored at each vertex. Note that the simple way
           of creating a graph does not allow you to assign data, you must
           do it manually."""

        def __init__ (self, vertices=[], edges=[]):
                """Constructor"""
                self.vertices = {}
                self.edges = {}

                for v in vertices:
                        self.add_vertex (v)

                for e in edges:
                        (v1, v2) = e
                        self.add_edge (v1, v2)

        def __get_key (self, v1, v2):
                """Create a key for this set of nodes"""
                key = [v1, v2]
                key.sort ()
                return tuple (key)

        def add_vertex (self, label, data=0, raw_obj=None):
                """Add a vertex to the graph"""
                if self.has_vertex (label):
                        raise ValueError

                # Create vertex and set the value
                self.vertices[label] = Vertex(label, data, raw_obj=raw_obj)

        def add_edge (self, v1, v2, color='black', label=''):
                """Add an edge from v1 to v2 (undirected)"""

                key = self.__get_key (v1, v2)

                if self.has_vertex (v1) and self.has_vertex (v2):
                        self.edges[key] = Edge (color, label)
                else:
                        raise ValueError

        def has_edge (self, v1, v2):
                """Check if an edge exists in this graph"""

                key = self.__get_key (v1, v2)

                if key in self.edges.keys():
                        return True

                return False

        def has_vertex (self, label):
                """Check if a vertex exists in this graph"""
                return label in self.vertices.keys ()

        def get_children (self, v):
                """Get a list of all children of v. In this graph, that
                   means any adjacent nodes."""
                children = set()

                # Check that the vertex exists
                if not self.has_vertex (v):
                        raise ValueError

                for k in self.edges.keys ():
                        if v in k:
                                for i in k:
                                        if i != v:
                                                children.add (i)

                return list(children)

        def set_vertex_value (self, v, value):
                """Store some data at a vertex"""
                if not self.has_vertex (v):
                        raise ValueError

                self.vertices[v].value = value

        def get_vertex_value (self, v):
                """Get the data at a vertex"""
                if not self.has_vertex (v):
                        raise ValueError

                return self.vertices[v].value

        def set_vertex_shape (self, v, shape):
                if not self.has_vertex (v):
                        raise ValueError

                self.vertices[v].shape = shape
                
        def get_vertex_shape (self, v):
                if not self.has_vertex (v):
                        raise ValueError

                return self.vertices[v].shape

        def set_edge_color (self, v1, v2, color):
                if not self.has_edge (v1, v2):
                        raise ValueError

                key = self.__get_key (v1, v2)
                self.edges[key].color = color

        def get_edge_color (self, v1, v2):
                if not self.has_edge (v1, v2):
                        raise ValueError

                key = self.__get_key (v1, v2)
                return self.edges[key].color

        def get_edge_label (self, v1, v2):
                if not self.has_edge (v1, v2):
                        raise ValueError

                key = self.__get_key (v1, v2)
                return self.edges[key].label

        def set_edge_label (self, v1, v2, label):
                if not self.has_edge (v1, v2):
                        raise ValueError

                key = self.__get_key (v1, v2)
                self.edges[key].label = label


################################################################################
###                      Tests are below this point                          ###
################################################################################


def run_test (graph, vertex, should_be, suite, test_num):
        """Run a test, check the result. Print the status."""
        c = graph.get_children (vertex)
        c.sort ()

        if (c != should_be):
                print 'Failed test #%d in suite: %s' % (test_num, suite)

def run_tests_1 ():
        """Run the first set of tests."""
        SUITE = 'Tests 1'

        vertices = range (11)
        edges = ( [0, 1], [0, 9], [0, 8], [8, 10], [8, 7], [9, 6], [9, 5],
                        [6, 7], [1, 4], [1, 3], [1, 2], [2, 0] )

        g = Graph (vertices, edges)

        run_test (g, 0, [1, 2, 8, 9], SUITE, 0)
        run_test (g, 1, [0, 2, 3, 4], SUITE, 1)
        run_test (g, 2, [0, 1], SUITE, 2)
        run_test (g, 3, [1], SUITE, 3)
        run_test (g, 4, [1], SUITE, 4)
        run_test (g, 5, [9], SUITE ,5)
        run_test (g, 6, [7, 9], SUITE, 6)
        run_test (g, 7, [6, 8], SUITE, 7)
        run_test (g, 8, [0, 7, 10], SUITE, 8)
        run_test (g, 9, [0, 5, 6], SUITE, 9)
        run_test (g, 10, [8], SUITE, 10)

def run_tests_2 ():
        """Run the second set of tests."""
        SUITE = 'Tests 2'

        vertices = ['a', 'b', 'c', 'd', 'e']
        edges = ( ['a', 'b'], ['a', 'c'], ['a', 'd'], ['b', 'a'], ['b', 'c'],
                        ['b', 'e'], ['e', 'b'], ['e', 'c'], ['e', 'd'],
                        ['d', 'a'], ['d', 'c'], ['d', 'e'], ['c', 'a'],
                        ['c', 'b'], ['c', 'd'], ['c', 'e'] )

        g = Graph (vertices, edges)

        run_test (g, 'a', ['b', 'c', 'd'], SUITE, 0)
        run_test (g, 'b', ['a', 'c', 'e'], SUITE, 1)
        run_test (g, 'c', ['a', 'b', 'd', 'e'], SUITE, 2)
        run_test (g, 'd', ['a', 'c', 'e'], SUITE, 3)
        run_test (g, 'e', ['b', 'c', 'd'], SUITE, 4)

def main ():
        """Main function."""

        print 'Running self tests. Any failures will be reported below'

        run_tests_1 ()
        run_tests_2 ()


if __name__ == '__main__':
        main ()