Merge

1 files changed, 41 insertions, 13 deletions

diff --git a/examples/source.py b/examples/source.py
index 6bfdcb0..adaeeb3 100644
--- a/examples/source.py
+++ b/examples/source.py
@@ -1,15 +1,36 @@
 """Test client that connects and sends infinite data."""
 
+import argparse
 import sys
 
 from asyncio import *
 
+
+ARGS = argparse.ArgumentParser(description="TCP data sink example.")
+ARGS.add_argument(
+    '--iocp', action='store_true', dest='iocp',
+    default=False, help='Use IOCP event loop (Windows only)')
+ARGS.add_argument(
+    '--stop', action='store_true', dest='stop',
+    default=False, help='Stop the server by sending it b"stop" as data')
+ARGS.add_argument(
+    '--host', action='store', dest='host',
+    default='127.0.0.1', help='Host name')
+ARGS.add_argument(
+    '--port', action='store', dest='port',
+    default=1111, type=int, help='Port number')
+ARGS.add_argument(
+    '--size', action='store', dest='size',
+    default=16*1024, type=int, help='Data size')
+
+args = None
+
+
 def dprint(*args):
     print('source:', *args, file=sys.stderr)
 
-class Client(Protocol):
 
-    data = b'x'*16*1024
+class Client(Protocol):
 
     def connection_made(self, tr):
         dprint('connecting to', tr.get_extra_info('peername'))
@@ -18,41 +39,48 @@ class Client(Protocol):
         self.lost = False
         self.loop = get_event_loop()
         self.waiter = Future()
-        if '--stop' in sys.argv[1:]:
+        if args.stop:
             self.tr.write(b'stop')
             self.tr.close()
         else:
-            self.write_some_data()
+            data = b'x' * args.size
+            self.write_some_data(data)
 
-    def write_some_data(self):
+    def write_some_data(self, data):
         if self.lost:
             dprint('lost already')
             return
-        dprint('writing', len(self.data), 'bytes')
-        self.tr.write(self.data)
-        self.loop.call_soon(self.write_some_data)
+        dprint('writing', len(data), 'bytes')
+        self.tr.write(data)
+        self.loop.call_soon(self.write_some_data, data)
 
     def connection_lost(self, exc):
         dprint('lost connection', repr(exc))
         self.lost = True
         self.waiter.set_result(None)
 
+
 @coroutine
-def start(loop):
-    tr, pr = yield from loop.create_connection(Client, '127.0.0.1', 1111)
+def start(loop, host, port):
+    tr, pr = yield from loop.create_connection(Client, host, port)
     dprint('tr =', tr)
     dprint('pr =', pr)
     res = yield from pr.waiter
     return res
 
+
 def main():
-    if '--iocp' in sys.argv:
+    global args
+    args = ARGS.parse_args()
+    if args.iocp:
         from asyncio.windows_events import ProactorEventLoop
         loop = ProactorEventLoop()
         set_event_loop(loop)
-    loop = get_event_loop()
-    loop.run_until_complete(start(loop))
+    else:
+        loop = get_event_loop()
+    loop.run_until_complete(start(loop, args.host, args.port))
     loop.close()
 
+
 if __name__ == '__main__':
     main()