diff --git a/greentulip/socket.py b/greentulip/socket.py index cf5c851..8e0d18e 100644 --- a/greentulip/socket.py +++ b/greentulip/socket.py @@ -4,6 +4,14 @@ ## +"""Greensocket (non-blocking) for Tulip. + +Use ``greentulip.socket`` in the same way as you would use stdlib's +``socket.socket`` in ``greentulip.task`` tasks or coroutines invoked +from them. +""" + + import tulip from socket import * @@ -36,34 +44,73 @@ class socket: def proto(self): return self._sock.proto - def setblocking(flag): - assert not flag, 'greenlet.socket does not support blocking mode' + def _proxy(attr): + def proxy(self, *args, **kwargs): + meth = getattr(self._sock, attr) + return meth(*args, **kwargs) + proxy.__name__ = attr + proxy.__qualname__ = attr + proxy.__doc__ = getattr(getattr(std_socket, attr), '__doc__', None) + return proxy + + def _copydoc(func): + func.__doc__ = getattr(getattr(std_socket, func.__name__), '__doc__', None) + return func + + @_copydoc + def setblocking(self, flag): + if flag: + raise error('greentulip.socket does not support blocking mode') + + @_copydoc def recv(self, nbytes): fut = tulip.get_event_loop().sock_recv(self._sock, nbytes) - return yield_from(fut) + yield_from(fut) + return fut.result() + @_copydoc def connect(self, addr): fut = tulip.get_event_loop().sock_connect(self._sock, addr) - return yield_from(fut) + yield_from(fut) + return fut.result() + @_copydoc def sendall(self, data, flags=0): - assert not flags - fut = tulip.get_event_loop().sock_sendall(self._sock, data) - return yield_from(fut) - - def send(self, data, flags=0): assert not flags fut = tulip.get_event_loop().sock_sendall(self._sock, data) yield_from(fut) + return fut.result() + + @_copydoc + def send(self, data, flags=0): + self.sendall(data, flags) return len(data) + @_copydoc def accept(self): fut = tulip.get_event_loop().sock_accept(self._sock) - return yield_from(fut) + yield_from(fut) + sock, addr = fut.result() + return self.__class__.from_socket(sock), addr - def close(self): - return self._sock.close() + @_copydoc + def makefile(self, *args, **kwargs): + raise NotImplementedError + + bind = _proxy('bind') + listen = _proxy('listen') + getsockname = _proxy('getsockname') + getpeername = _proxy('getpeername') + gettimeout = _proxy('gettimeout') + getsockopt = _proxy('getsockopt') + setsockopt = _proxy('setsockopt') + fileno = _proxy('fileno') + detach = _proxy('detach') + close = _proxy('close') + shutdown = _proxy('shutdown') + + del _copydoc, _proxy def create_connection(address:tuple, timeout=None): diff --git a/tests/test_socket.py b/tests/test_socket.py new file mode 100644 index 0000000..cfba4a3 --- /dev/null +++ b/tests/test_socket.py @@ -0,0 +1,109 @@ +## +# Copyright (c) 2013 Yury Selivanov +# License: Apache 2.0 +## + + +import greentulip +import greentulip.socket as greensocket + +import tulip +import unittest + + +class SocketTests(unittest.TestCase): + def setUp(self): + tulip.set_event_loop_policy(greentulip.GreenEventLoopPolicy()) + self.loop = tulip.new_event_loop() + tulip.set_event_loop(self.loop) + + def tearDown(self): + self.loop.close() + tulip.set_event_loop_policy(None) + + def test_socket_docs(self): + self.assertTrue('accept connections' in greensocket.socket.listen.__doc__) + self.assertTrue('Receive' in greensocket.socket.recv.__doc__) + + def test_socket_setblocking(self): + sock = greensocket.socket() + self.assertEquals(sock.gettimeout(), 0) + with self.assertRaisesRegex(greensocket.error, 'does not support blocking mode'): + sock.setblocking(True) + + def test_socket_echo(self): + import socket as std_socket + import threading + import time + + check = 0 + ev = threading.Event() + + def server(sock_factory): + socket = sock_factory() + socket.bind(('127.0.0.1', 0)) + + assert socket.fileno() is not None + + nonlocal addr + addr = socket.getsockname() + socket.listen(1) + + ev.set() + + sock, client_addrs = socket.accept() + assert isinstance(sock, sock_factory) + + data = b'' + while not data.endswith(b'\r'): + data += sock.recv(1024) + + sock.sendall(data) + + ev.wait() + ev.clear() + + sock.close() + socket.close() + + def client(sock_factory): + ev.wait() + ev.clear() + time.sleep(0.1) + + assert addr + sock = sock_factory() + sock.connect(addr) + + data = b'hello greenlets\r' + sock.sendall(data) + + rep = b'' + while not rep.endswith(b'\r'): + rep += sock.recv(1024) + + self.assertEqual(data, rep) + ev.set() + + nonlocal check + check += 1 + + sock.close() + + addr = None + ev.clear() + thread = threading.Thread(target=client, args=(std_socket.socket,)) + thread.setDaemon(True) + thread.start() + self.loop.run_until_complete(greentulip.task(server)(greensocket.socket)) + thread.join(1) + self.assertEqual(check, 1) + + addr = None + ev.clear() + thread = threading.Thread(target=server, args=(std_socket.socket,)) + thread.setDaemon(True) + thread.start() + self.loop.run_until_complete(greentulip.task(client)(greensocket.socket)) + thread.join(1) + self.assertEqual(check, 2)