From 8e8d43b0413437273d6c7d7f62ef19ba699c4313 Mon Sep 17 00:00:00 2001
From: Roman Zeyde <roman.zeyde@gmail.com>
Date: Fri, 9 Jan 2015 19:50:04 +0200
Subject: [PATCH] cli: add zip compression option

---
 amodem-cli     | 44 +++++++++++++++++++++++++++++++++++++++-----
 amodem/recv.py |  1 +
 2 files changed, 40 insertions(+), 5 deletions(-)

diff --git a/amodem-cli b/amodem-cli
index 52a1a7a..6b91778 100755
--- a/amodem-cli
+++ b/amodem-cli
@@ -1,7 +1,12 @@
 #!/usr/bin/env python
 # PYTHON_ARGCOMPLETE_OK
 
+import os
 import sys
+import zlib
+import logging
+import argparse
+
 if sys.version_info.major == 2:
     _stdin = sys.stdin
     _stdout = sys.stdout
@@ -9,23 +14,47 @@ else:
     _stdin = sys.stdin.buffer
     _stdout = sys.stdout.buffer
 
-import argparse
 try:
     import argcomplete
 except ImportError:
     argcomplete = None
 
-import logging
 log = logging.getLogger('__name__')
 
 from amodem import recv, send, calib, audio
 from amodem.config import bitrates
 
-import os
 bitrate = os.environ.get('BITRATE', 1)
 config = bitrates.get(int(bitrate))
 
 
+class Compressor(object):
+    def __init__(self, stream):
+        self.obj = zlib.compressobj()
+        self.stream = stream
+
+    def read(self, size):
+        data = self.stream.read(size)
+        if data:
+            data = self.obj.compress(data)
+        elif self.obj:
+            data = self.obj.flush()
+            self.obj = None
+        return data
+
+
+class Decompressor(object):
+    def __init__(self, stream):
+        self.obj = zlib.decompressobj()
+        self.stream = stream
+
+    def write(self, data):
+        self.stream.write(self.obj.decompress(bytes(data)))
+
+    def flush(self):
+        self.stream.write(self.obj.flush())
+
+
 def FileType(mode, audio_interface=None):
     def opener(fname):
         assert 'r' in mode or 'w' in mode
@@ -60,6 +89,9 @@ def main():
     p = argparse.ArgumentParser(description=description)
     subparsers = p.add_subparsers()
 
+    def wrap(cls, stream, enable):
+        return cls(stream) if enable else stream
+
     # Modulator
     sender = subparsers.add_parser(
         'send', help='modulate binary data into audio signal.')
@@ -73,7 +105,7 @@ def main():
 
     sender.set_defaults(
         main=lambda config, args: send.main(
-            config, src=args.src, dst=args.dst
+            config, src=wrap(Compressor, args.src, args.zip), dst=args.dst
         ),
         calib=lambda config, args: calib.send(
             config, dst=args.dst
@@ -97,7 +129,8 @@ def main():
         help='plot results using pylab module')
     receiver.set_defaults(
         main=lambda config, args: recv.main(
-            config, src=args.src, dst=args.dst, pylab=args.pylab
+            config, src=args.src, dst=wrap(Decompressor, args.dst, args.zip),
+            pylab=args.pylab
         ),
         calib=lambda config, args: calib.recv(
             config, src=args.src, verbose=args.verbose
@@ -107,6 +140,7 @@ def main():
     )
 
     for sub in subparsers.choices.values():
+        sub.add_argument('-z', '--zip', default=False, action='store_true')
         g = sub.add_mutually_exclusive_group()
         g.add_argument('-v', '--verbose', default=0, action='count')
         g.add_argument('-q', '--quiet', default=False, action='store_true')
diff --git a/amodem/recv.py b/amodem/recv.py
index 5147e45..f52ec7d 100644
--- a/amodem/recv.py
+++ b/amodem/recv.py
@@ -232,5 +232,6 @@ def main(config, src, dst, pylab=None):
     except Exception:
         log.exception('Decoding failed')
 
+    dst.flush()
     receiver.report()
     return success