python tags source/sink looper

2 files changed, 28 insertions, 14 deletions

diff --git a/tests/block_test.py b/tests/block_test.py
index 9421b66..ff44a74 100644
--- a/tests/block_test.py
+++ b/tests/block_test.py
@@ -44,13 +44,15 @@ class BlockTest(unittest.TestCase):
         tb = None
         self.assertEqual(sink.get_vector(), (1, 5j, 9, 13j, 17))
 
-    def test_tag_source(self):
+    def test_tag_source_sink(self):
+        values = (0, 'hello', 4.2, True, None, [2, 3, 4], (9, 8, 7), 1j, {2:'d'})
         tb = gras.TopBlock()
-        src = TagSource([1, 2, 3])
-        sink = NullSink(numpy.uint8)
+        src = TagSource(values)
+        sink = TagSink()
         tb.connect(src, sink)
         tb.run()
         tb = None
+        self.assertEqual(sink.get_values(), values)
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/tests/demo_blocks.py b/tests/demo_blocks.py
index 84c8e8e..9a7a02f 100644
--- a/tests/demo_blocks.py
+++ b/tests/demo_blocks.py
@@ -16,10 +16,7 @@ class NullSink(gras.Block):
         gras.Block.__init__(self, 'NullSink', in_sig=[in_sig])
 
     def work(self, ins, outs):
-        print 'consume',len(ins[0])
         self.consume(0, len(ins[0]))
-        #FIXME should not need == 0 part!!
-        if len(ins[0]) == 0: self.mark_done()
 
 class VectorSource(gras.Block):
     def __init__(self, out_sig, vec):
@@ -62,20 +59,35 @@ class Add2X(gras.Block):
         self.produce(0, nitems)
 
 class TagSource(gras.Block):
-    def __init__(self, tags):
+    def __init__(self, values):
         gras.Block.__init__(self,
             name = "TagSource",
             out_sig = [numpy.uint8],
         )
-        self._tags = tags
+        self._values = values
 
     def work(self, ins, outs):
         offset = self.nitems_written(0)
-        self.post_output_tag(0, gras.Tag(offset, 'key', self._tags[0]))
+        self.post_output_tag(0, gras.Tag(offset, 'key', self._values[0]))
         self.produce(0, len(outs[0]))
-        self._tags = self._tags[1:]
-        print 'produce', len(outs[0])
-        print 'self._tags', self._tags
-        if not self._tags:
+        self._values = self._values[1:]
+        if not self._values:
             self.mark_done()
-            print 'done!'
+
+class TagSink(gras.Block):
+    def __init__(self):
+        gras.Block.__init__(self,
+            name = "TagSink",
+            in_sig = [numpy.uint8],
+        )
+        self._values = list()
+
+    def get_values(self):
+        return tuple(self._values)
+
+    def work(self, ins, outs):
+        max_read = self.nitems_read(0) + len(ins[0])
+        for tag in self.get_input_tags(0):
+            if tag.offset < max_read:
+                self._values.append(tag.value)
+        self.consume(0, len(ins[0]))