Merge pull request #13 from dirk-thomas/custom_highlight_color

Add optional argument for custom highlight color

gvanim/action.py

@@ -29,11 +29,12 @@ def __call__( self, steps ):
 		steps[ -1 ].V.add( self.v )
 
 class HighlightNode( object ):
-	def __init__( self, v ):
+	def __init__( self, v, color = 'red' ):
 		self.v = v
+		self.color = color
 	def __call__( self, steps ):
 		steps[ -1 ].V.add( self.v )
-		steps[ -1 ].hV.add( self.v )
+		steps[ -1 ].hV[ self.v ] = self.color
 
 class LabelNode( object ):
 	def __init__( self, v, label ):
@@ -58,14 +59,19 @@ def __init__( self, v ):
 		self.v = v
 	def __call__( self, steps ):
 		steps[ -1 ].V.discard( self.v )
-		steps[ -1 ].hV.discard( self.v )
+		try:
+			del steps[ -1 ].hV[ self.v ]
+		except KeyError:
+			pass
 		try:
 			del steps[ -1 ].L[ self.v ]
 		except KeyError:
 			pass
 		dE = set( e for e in steps[ -1 ].E if self.v in e )
 		steps[ -1 ].E -= dE
-		steps[ -1 ].hE -= dE
+		for e in list(steps[ -1 ].hE.keys()):
+			if self.v in e:
+				del steps[ -1 ].hE[ e ]
 
 class AddEdge( object ):
 	def __init__( self, u, v ):
@@ -77,20 +83,24 @@ def __call__( self, steps ):
 		steps[ -1 ].E.add( ( self.u, self.v ) )
 
 class HighlightEdge( object ):
-	def __init__( self, u, v ):
+	def __init__( self, u, v, color = 'red' ):
 		self.u = u
 		self.v = v
+		self.color = color
 	def __call__( self, steps ):
 		steps[ -1 ].V.add( self.u )
 		steps[ -1 ].V.add( self.v )
 		steps[ -1 ].E.add( ( self.u, self.v ) )
-		steps[ -1 ].hE.add( ( self.u, self.v ) )
+		steps[ -1 ].hE[ ( self.u, self.v ) ] = self.color
 
 class RemoveEdge( object ):
 	def __init__( self, u, v ):
 		self.u = u
 		self.v = v
 	def __call__( self, steps ):
 		steps[ -1 ].E.discard( ( self.u, self.v ) )
-		steps[ -1 ].hE.discard( ( self.u, self.v ) )
+		try:
+			del steps[ -1 ].hE[ ( self.u, self.v ) ]
+		except KeyError:
+			pass
 

gvanim/animation.py

@@ -34,8 +34,8 @@ def __init__( self, step = None ):
 			self.V = set()
 			self.E = set()
 			self.L = dict()
-		self.hV = set()
-		self.hE = set()
+		self.hV = dict()
+		self.hE = dict()
 
 	def node_format( self, v ):
 		fmt = []
@@ -44,15 +44,15 @@ def node_format( self, v ):
 		except KeyError:
 			pass
 		if v in self.hV:
-			fmt.append( 'color=red' )
+			fmt.append( 'color={}'.format( self.hV[ v ] ) )
 		elif v not in self.V:
 			fmt.append( 'style=invis' )
 		if fmt: return '[{}]'.format( ','.join( fmt ) )
 		return ''
 
 	def edge_format( self, e ):
 		if e in self.hE:
-			return '[color=red]'
+			return '[color={}]'.format( self.hE[ e ] )
 		elif e in self.E:
 			return ''
 		return '[style=invis]'
@@ -71,8 +71,8 @@ def next_step( self, clean = False ):
 	def add_node( self, v ):
 		self._actions.append( action.AddNode( v ) )
 
-	def highlight_node( self, v ):
-		self._actions.append( action.HighlightNode( v ) )
+	def highlight_node( self, v, color = 'red' ):
+		self._actions.append( action.HighlightNode( v, color = color ) )
 
 	def label_node( self, v, label ):
 		self._actions.append( action.LabelNode( v, label ) )
@@ -86,8 +86,8 @@ def remove_node( self, v ):
 	def add_edge( self, u, v ):
 		self._actions.append( action.AddEdge( u, v ) )
 
-	def highlight_edge( self, u, v ):
-		self._actions.append( action.HighlightEdge( u, v ) )
+	def highlight_edge( self, u, v, color = 'red' ):
+		self._actions.append( action.HighlightEdge( u, v, color = color ) )
 
 	def remove_edge( self, u, v ):
 		self._actions.append( action.RemoveEdge( u, v ) )