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 ()