changed add_state_hook to set_state_hook

bniemczyk · Jan 27, 2014 · 37c9645 · 37c9645
1 parent 1005b68
commit 37c9645
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/automata/automata.py b/automata/automata.py
@@ -34,8 +34,12 @@ def _choose(self, arglist):
       else:
         rv = self.choose(rv, i)
 
-  def add_state_hook(self, state, hook):
-    self._state_hooks[state] = hook
+  def set_state_hook(self, state, hook):
+    if hook == None:
+      if state in self._state_hooks:
+        del self._state_hooks[state]
+    else:
+      self._state_hooks[state] = hook
 
   def transitions_to(self, dst):
     '''

diff --git a/tests/automata_test.py b/tests/automata_test.py
@@ -62,6 +62,6 @@ def _h():
       _h.count += 1
       return _h.count
     _h.count = 0
-    nfa.add_state_hook(0, _h)
+    nfa.set_state_hook(0, _h)
     nfa.execute('aaaba')
     self.assertEqual(_h.count, 3)