working on tags python support

2 files changed, 59 insertions, 16 deletions

diff --git a/tests/block_test.py b/tests/block_test.py
index 23322a4..6d922a9 100644
--- a/tests/block_test.py
+++ b/tests/block_test.py
@@ -21,7 +21,7 @@ class BlockTest(unittest.TestCase):
         tb = gras.TopBlock()
         src0 = VectorSource(numpy.float32, [1, 3, 5, 7, 9])
         src1 = VectorSource(numpy.float32, [0, 2, 4, 6, 8])
-        adder = add_2x(numpy.float32)
+        adder = Add2X(numpy.float32)
         sink = VectorSink(numpy.float32)
         tb.connect((src0, 0), (adder, 0))
         tb.connect((src1, 0), (adder, 1))
@@ -30,5 +30,26 @@ class BlockTest(unittest.TestCase):
         tb = None
         self.assertEqual(sink.get_vector(), (1, 5, 9, 13, 17))
 
+    def test_add_fc32(self):
+        tb = gras.TopBlock()
+        src0 = VectorSource(numpy.complex64, [1, 3j, 5, 7j, 9])
+        src1 = VectorSource(numpy.complex64, [0, 2j, 4, 6j, 8])
+        adder = Add2X(numpy.complex64)
+        sink = VectorSink(numpy.complex64)
+        tb.connect((src0, 0), (adder, 0))
+        tb.connect((src1, 0), (adder, 1))
+        tb.connect(adder, sink)
+        tb.run()
+        tb = None
+        self.assertEqual(sink.get_vector(), (1, 5j, 9, 13j, 17))
+
+    def test_tag_source(self):
+        tb = gras.TopBlock()
+        src = TagSource([1, 2, 3])
+        sink = NullSink(numpy.uint8)
+        tb.connect(src, sink)
+        tb.run()
+        tb = None
+
 if __name__ == '__main__':
     unittest.main()
diff --git a/tests/demo_blocks.py b/tests/demo_blocks.py
index 3a5d5cb..7187ec7 100644
--- a/tests/demo_blocks.py
+++ b/tests/demo_blocks.py
@@ -1,21 +1,6 @@
 import gras
 import numpy
 
-class add_2x(gras.Block):
-    def __init__(self, sig):
-        gras.Block.__init__(self,
-            name = "add 2x",
-            in_sig = [sig, sig],
-            out_sig = [sig],
-        )
-
-    def work(self, ins, outs):
-        nitems = min(*map(len, (ins[0], ins[1], outs[0])))
-        outs[0][:nitems] = ins[0][:nitems] + ins[1][:nitems]
-        self.consume(0, nitems)
-        self.consume(1, nitems)
-        self.produce(0, nitems)
-
 class NullSource(gras.Block):
     def __init__(self, out_sig):
         gras.Block.__init__(self, 'NullSource', out_sig=[out_sig])
@@ -29,7 +14,10 @@ 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):
@@ -55,3 +43,37 @@ class VectorSink(gras.Block):
     def work(self, ins, outs):
         self._vec.extend(ins[0].copy())
         self.consume(0, len(ins[0]))
+
+class Add2X(gras.Block):
+    def __init__(self, sig):
+        gras.Block.__init__(self,
+            name = "Add2X",
+            in_sig = [sig, sig],
+            out_sig = [sig],
+        )
+
+    def work(self, ins, outs):
+        nitems = min(*map(len, (ins[0], ins[1], outs[0])))
+        outs[0][:nitems] = ins[0][:nitems] + ins[1][:nitems]
+        self.consume(0, nitems)
+        self.consume(1, nitems)
+        self.produce(0, nitems)
+
+class TagSource(gras.Block):
+    def __init__(self, tags):
+        gras.Block.__init__(self,
+            name = "TagSource",
+            out_sig = [numpy.uint8],
+        )
+        self._tags = tags
+
+    def work(self, ins, outs):
+        offset = self.nitems_written(0)
+        self.post_output_tag(0, gras.Tag(offset, 'key', self._tags[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.mark_done()
+            print 'done!'