Reader: fix iteration logic.

amodem/send.py

@@ -13,6 +13,7 @@ import wave
 import common
 import config
 import sigproc
+import stream
 
 modem = sigproc.MODEM(config)
 
@@ -64,24 +65,6 @@ def modulate(fd, bits):
             break
 
 
-class Reader(object):
-    def __init__(self, fd, size):
-        self.fd = fd
-        self.size = size
-        self.total = 0
-
-    def __next__(self):
-        block = self.fd.read(self.size)
-        if block:
-            self.total += len(block)
-            return block
-        else:
-            raise StopIteration()
-
-    def __iter__(self):
-        return self
-
-
 def main(args):
     import ecc
     log.info('Running MODEM @ {:.1f} kbps'.format(modem.modem_bps / 1e3))
@@ -96,7 +79,7 @@ def main(args):
     log.info('%.3f seconds of training audio',
              training_size / wave.bytes_per_second)
 
-    reader = Reader(args.input, 64 << 10)
+    reader = stream.Reader(args.input, 64 << 10)
     data = itertools.chain.from_iterable(reader)
     encoded = itertools.chain.from_iterable(ecc.encode(data))
     modulate(args.output, bits=common.to_bits(encoded))

amodem/stream.py

@@ -12,21 +12,30 @@ class Reader(object):
     WAIT = 0.1
     TIMEOUT = 2.0
 
-    def __init__(self, fd):
+    def __init__(self, fd, bufsize=None, eof=False):
         self.fd = fd
         self.check = None
+        self.total = 0
+        self.bufsize = bufsize if (bufsize is not None) else self.BUFSIZE
+        self.eof = eof
 
     def __iter__(self):
         return self
 
     def __next__(self):
+        return self.next()
+
+    def next(self):
         block = bytearray()
         finish_time = time.time() + self.TIMEOUT
         while time.time() <= finish_time:
             left = self.BUFSIZE - len(block)
             data = self.fd.read(left)
             if data:
+                self.total += len(data)
                 block.extend(data)
+            elif self.eof:  # handle EOF condition by stopping iteration
+                raise StopIteration()
 
             if len(block) == self.BUFSIZE:
                 values = common.loads(str(block))