Improved argument handling.

2 files changed, 14 insertions, 7 deletions

diff --git a/examples/sink.py b/examples/sink.py
index 119af6e..c94574a 100644
--- a/examples/sink.py
+++ b/examples/sink.py
@@ -15,8 +15,12 @@ ARGS.add_argument(
 ARGS.add_argument(
     '--port', action='store', dest='port',
     default=1111, type=int, help='Port number')
+ARGS.add_argument(
+    '--maxsize', action='store', dest='maxsize',
+    default=16*1024*1024, type=int, help='Max total data size')
 
 server = None
+args = None
 
 
 def dprint(*args):
@@ -39,7 +43,7 @@ class Service(Protocol):
             return
         self.total += len(data)
         dprint('received', len(data), 'bytes; total', self.total)
-        if self.total > 1e6:
+        if self.total > args.maxsize:
             dprint('closing due to too much data')
             self.tr.close()
 
@@ -56,6 +60,7 @@ def start(loop, host, port):
 
 
 def main():
+    global args
     args = ARGS.parse_args()
     if args.iocp:
         from tulip.windows_events import ProactorEventLoop
diff --git a/examples/source.py b/examples/source.py
index 1515dd1..933739b 100644
--- a/examples/source.py
+++ b/examples/source.py
@@ -45,17 +45,19 @@ class Client(Protocol):
             self.tr.write(b'stop')
             self.tr.close()
         else:
-            data = b'x' * args.size
-            self.write_some_data(data)
+            self.data = b'x'*args.size
+            self.write_some_data()
 
-    def write_some_data(self, data):
+    def write_some_data(self):
         if self.lost:
             dprint('lost already')
             return
-        self.total += len(data)
-        dprint('writing', len(data), 'bytes; total', self.total)
+        data = self.data
+        size = len(data)
+        self.total += size
+        dprint('writing', size, 'bytes; total', self.total)
         self.tr.write(data)
-        self.loop.call_soon(self.write_some_data, data)
+        self.loop.call_soon(self.write_some_data)
 
     def connection_lost(self, exc):
         dprint('lost connection', repr(exc))