Also accept buffer in Connection.send and .sendall

Not at all sure what i am doing, but according to
http://hg.python.org/cpython/rev/9e718d8f71e8/, buffers are a accepted
type to pass to sendall in Python 2.6.

3 files changed, 51 insertions, 4 deletions

diff --git a/ChangeLog b/ChangeLog
index 164c9ad..53fb1df 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -1,3 +1,8 @@
+2014-04-19  Markus Unterwaditzer <markus@unterwaditzer.net>
+
+	* OpenSSL/SSL.py: ``Connection.send`` and ``Connection.sendall``
+	  now also accept the ``buffer`` type as data.
+
 2014-03-30  Fedor Brunner  <fedor.brunner@azet.sk>
 
 	* OpenSSL/SSL.py: Add ``get_finished``, ``get_peer_finished``
diff --git a/OpenSSL/SSL.py b/OpenSSL/SSL.py
index fbb18f0..593d89f 100644
--- a/OpenSSL/SSL.py
+++ b/OpenSSL/SSL.py
@@ -24,6 +24,12 @@ except NameError:
     class _memoryview(object):
         pass
 
+try:
+    _buffer = buffer
+except NameError:
+    class _buffer(object):
+        pass
+
 OPENSSL_VERSION_NUMBER = _lib.OPENSSL_VERSION_NUMBER
 SSLEAY_VERSION = _lib.SSLEAY_VERSION
 SSLEAY_CFLAGS = _lib.SSLEAY_CFLAGS
@@ -940,15 +946,17 @@ class Connection(object):
         WantWrite or WantX509Lookup exceptions on this, you have to call the
         method again with the SAME buffer.
 
-        :param buf: The string to send
+        :param buf: The string, buffer or memoryview to send
         :param flags: (optional) Included for compatibility with the socket
                       API, the value is ignored
         :return: The number of bytes written
         """
         if isinstance(buf, _memoryview):
             buf = buf.tobytes()
+        if isinstance(buf, _buffer):
+            buf = str(buf)
         if not isinstance(buf, bytes):
-            raise TypeError("data must be a byte string")
+            raise TypeError("data must be a memoryview, buffer or byte string")
 
         result = _lib.SSL_write(self._ssl, buf, len(buf))
         self._raise_ssl_error(self._ssl, result)
@@ -962,15 +970,17 @@ class Connection(object):
         all data is sent. If an error occurs, it's impossible to tell how much
         data has been sent.
 
-        :param buf: The string to send
+        :param buf: The string, buffer or memoryview to send
         :param flags: (optional) Included for compatibility with the socket
                       API, the value is ignored
         :return: The number of bytes written
         """
         if isinstance(buf, _memoryview):
             buf = buf.tobytes()
+        if isinstance(buf, _buffer):
+            buf = str(buf)
         if not isinstance(buf, bytes):
-            raise TypeError("buf must be a byte string")
+            raise TypeError("buf must be a memoryview, buffer or byte string")
 
         left_to_send = len(buf)
         total_sent = 0
diff --git a/OpenSSL/test/test_ssl.py b/OpenSSL/test/test_ssl.py
index bfe3114..abc3bbe 100644
--- a/OpenSSL/test/test_ssl.py
+++ b/OpenSSL/test/test_ssl.py
@@ -2136,6 +2136,23 @@ class ConnectionSendTests(TestCase, _LoopbackMixin):
             self.assertEquals(client.recv(2), b('xy'))
 
 
+    try:
+        buffer
+    except NameError:
+        "cannot test sending buffer without buffer"
+    else:
+        def test_short_buffer(self):
+            """
+            When passed a buffer containing a small number of bytes,
+            :py:obj:`Connection.send` transmits all of them and returns the number of
+            bytes sent.
+            """
+            server, client = self._loopback()
+            count = server.send(buffer(b('xy')))
+            self.assertEquals(count, 2)
+            self.assertEquals(client.recv(2), b('xy'))
+
+
 
 class ConnectionSendallTests(TestCase, _LoopbackMixin):
     """
@@ -2179,6 +2196,21 @@ class ConnectionSendallTests(TestCase, _LoopbackMixin):
             self.assertEquals(client.recv(1), b('x'))
 
 
+    try:
+        buffer
+    except NameError:
+        "cannot test sending buffers without buffers"
+    else:
+        def test_short_buffers(self):
+            """
+            When passed a buffer containing a small number of bytes,
+            :py:obj:`Connection.sendall` transmits all of them.
+            """
+            server, client = self._loopback()
+            server.sendall(buffer(b('x')))
+            self.assertEquals(client.recv(1), b('x'))
+
+
     def test_long(self):
         """
         :py:obj:`Connection.sendall` transmits all of the bytes in the string passed to