Commits

Miki Tebeka committed 1983fa8

initial implementation

Comments (0)

Files changed (3)

+'''iocp tries to mimic Go's io.Copy function (http://golang.org/pkg/io/#Copy)
+which copies from everything that is an io.Reader to anything that is an
+io.Writer).
+
+The Python world of things is not that organized, so iocp.copy will try to do a
+best guess about which are the right function to call when reading and writing.
+'''
+
+READ_FNS = [
+    'read',  # file, StringIO ...
+    'recv',  # socket
+]
+
+WRITE_FNS = [
+    'write',  # file, StringIO, ...
+    'send',  # socket
+    'update',  # hashlib ...
+]
+
+BUFSIZE = 32*1024
+
+
+def find_fn(obj, names, name):
+    '''Find function in from list of names, raise ValueError not found.'''
+    for name in names:
+        fn = getattr(obj, name, None)
+        if fn:
+            return fn
+
+    raise ValueError('cannot find {} function for {}'.format(name, obj))
+
+
+def copy(src, dest, bufsize=BUFSIZE):
+    '''Copy from source to destination, return number of bytes read.
+
+    Since Python's "write" functions don't return number of bytes written, we
+    return number of bytes read as best estimation.
+    '''
+    write = find_fn(dest, WRITE_FNS, 'write')
+    read = find_fn(src, READ_FNS, 'read')
+
+    nwritten = 0
+    while True:
+            try:
+                buf = read(bufsize)
+            except EOFError:
+                buf = ''
+
+            if not buf:
+                break
+
+            nwritten += len(buf)
+            write(buf)
+
+    return nwritten
+#!/bin/bash
+
+nosetests -vd $@
+from os import urandom, SEEK_SET
+from tempfile import NamedTemporaryFile
+import hashlib
+from threading import Thread
+import socket
+
+import iocp
+
+
+def new_socket():
+    return socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+
+
+def socket_server():
+    sock = new_socket()
+    sock.bind(('localhost', 0))  # 0 will assign a random free port
+    port = sock.getsockname()[1]
+
+    def serve():
+        sock.listen(1)
+        while True:
+            conn, addr = sock.accept()
+            conn.send(data)
+            conn.close()
+
+    thr = Thread(target=serve)
+    thr.daemon = True
+    thr.start()
+
+    return port
+
+
+data_size = iocp.BUFSIZE * 3 + 17
+data = urandom(data_size)
+port = socket_server()
+
+
+def src_file():
+    src = NamedTemporaryFile()
+    src.write(data)
+    src.flush()
+    src.seek(0, SEEK_SET)
+
+    return src
+
+
+def src_socket():
+    sock = new_socket()
+    sock.connect(('localhost', port))
+
+    return sock
+
+
+def read_file(fo):
+    fo.seek(0, SEEK_SET)
+    return fo.read()
+
+
+def check_copy(src, dest, expected, readfn):
+    iocp.copy(src, dest)
+    expected = expected or data
+    received = readfn() if readfn else read_file(dest)
+    assert received == expected, 'bad copy {} -> {}'.format(src, dest)
+
+
+def test_copy():
+    hasher = hashlib.md5()
+    digest = hashlib.md5(data).hexdigest()
+
+    test_cases = [
+        (src_file(), NamedTemporaryFile(), None, None),
+        (src_file(), hasher, digest, lambda: hasher.hexdigest()),
+        (src_socket(), NamedTemporaryFile(), None, None),
+    ]
+
+    for src, dest, expected, readfn in test_cases:
+        yield check_copy, src, dest, expected, readfn