support '$' as trace filter terminator

2 files changed, 6 insertions, 1 deletions

diff --git a/lib/automata.py b/lib/automata.py
index 3220d81..71dcacc 100755
--- a/lib/automata.py
+++ b/lib/automata.py
@@ -85,6 +85,9 @@ class State:
         trace_filter -- list of lists. Each sub-list is a trace. Only traces matching one of the provided sub-lists are returned.
           E.g. trace_filter = [['init', 'foo'], ['init', 'bar']] will only return traces with init as first and foo or bar as second element.
         """
+
+        if trace_filter is not None and next(filter(lambda x: x == '$', map(lambda x: x[0], trace_filter)), None) is not None:
+            yield []
         if depth == 0:
             for trans in self.outgoing_transitions.values():
                 if trace_filter is not None and len(list(filter(lambda x: x == trans.name, map(lambda x: x[0], trace_filter)))) == 0:
@@ -100,7 +103,7 @@ class State:
                     yield [(trans,)]
         else:
             for trans in self.outgoing_transitions.values():
-                if trace_filter is not None and len(list(filter(lambda x: x == trans.name, map(lambda x: x[0], trace_filter)))) == 0:
+                if trace_filter is not None and next(filter(lambda x: x == trans.name, map(lambda x: x[0], trace_filter)), None) is None:
                     continue
                 if trace_filter is not None:
                     new_trace_filter = map(lambda x: x[1:], filter(lambda x: x[0] == trans.name, trace_filter))
diff --git a/test/test_pta.py b/test/test_pta.py
index 8ca41aa..dba67d7 100755
--- a/test/test_pta.py
+++ b/test/test_pta.py
@@ -162,6 +162,8 @@ class TestPTA(unittest.TestCase):
         pta.add_transition('IDLE', 'IDLE', 'set2')
         self.assertEqual(sorted(dfs_tran_to_name(pta.dfs(2, trace_filter=[['init', 'set1', 'set2'], ['init', 'set2', 'set1']]), False)),
             [['init', 'set1', 'set2'], ['init', 'set2', 'set1']])
+        self.assertEqual(sorted(dfs_tran_to_name(pta.dfs(2, trace_filter=[['init', 'set1', '$'], ['init', 'set2', '$']]), False)),
+            [['init', 'set1'], ['init', 'set2']])
 
     def test_dfs_accepting(self):
         pta = PTA(['IDLE', 'TX'], accepting_states = ['IDLE'])