Implement PTA DFS as generator

1 files changed, 8 insertions, 9 deletions

diff --git a/lib/automata.py b/lib/automata.py
index be91cd4..0e6cc28 100755
--- a/lib/automata.py
+++ b/lib/automata.py
@@ -15,15 +15,14 @@ class State:
 
     def dfs(self, depth):
         if depth == 0:
-            return [[trans.name] for trans in self.outgoing_transitions]
-
-        ret = []
-        for trans in self.outgoing_transitions:
-            for suffix in trans.destination.dfs(depth - 1):
-                new_suffix = [trans.name]
-                new_suffix.extend(suffix)
-                ret.append(new_suffix)
-        return ret
+            for trans in self.outgoing_transitions:
+                yield [trans.name]
+        else:
+            for trans in self.outgoing_transitions:
+                for suffix in trans.destination.dfs(depth - 1):
+                    new_suffix = [trans.name]
+                    new_suffix.extend(suffix)
+                    yield new_suffix
 
 class PTA:
     def __init__(self, state_names, parameters):